<commit_before>package authkit

// some constants
const (
	Facebook Provider = "facebook"
	Google            = "google"
	Github            = "github"
	Live              = "live"
	LinkedIn          = "linkedin"
)

// A ProviderRegistry contains all registerd providers.
type ProviderRegistry map[Provider]AuthRegistration

var (
	defaultBackends = ProviderRegistry{
		Facebook: AuthRegistration{
			Network:        Facebook,
			Scopes:         []string{"public_profile", "email"},
			AuthURL:        "https://www.facebook.com/dialog/oauth/",
			AccessTokenURL: "https://graph.facebook.com/oauth/access_token",
			AccessType:     "offline",
			UserinfoBase:   "https://graph.facebook.com/v3.0",
			UserinfoURLs:   []string{"/me?fields=email,name,gender,locale,picture"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].email",
			PathName:       "url[0].name",
			PathPicture:    "url[0].picture.data.url",
			PathCover:      "",
		},
		Google: AuthRegistration{
			Network:        Google,
			Scopes:         []string{"openid", "profile", "email", "https://www.googleapis.com/auth/plus.me"},
			AuthURL:        "https://accounts.google.com/o/oauth2/auth",
			AccessTokenURL: "https://accounts.google.com/o/oauth2/token",
			AccessType:     "offline",
			UserinfoBase:   "https://www.googleapis.com",
			UserinfoURLs:   []string{"/plus/v1/people/me"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].emails[0].value",
			PathName:       "url[0].displayName",
			PathPicture:    "url[0].image.url",
			PathCover:      "url[0].cover.coverPhoto.url",
		},
		Github: AuthRegistration{
			Network:        Github,
			Scopes:         []string{"user:email"},
			AuthURL:        "https://github.com/login/oauth/authorize",
			AccessTokenURL: "https://github.com/login/oauth/access_token",
			AccessType:     "offline",
			UserinfoBase:   "https://api.github.com",
			UserinfoURLs:   []string{"/user", "/user/emails"},
			PathID:         "url[0].login",
			PathEMail:      "url[1].data[0].email",
			PathName:       "url[0].name",
			PathPicture:    "url[0].avatar_url",
			PathCover:      "",
		},
		Live: AuthRegistration{
			Network:        Live,
			Scopes:         []string{"wl.signin", "wl.emails", "wl.basic"},
			AuthURL:        "https://login.live.com/oauth20_authorize.srf",
			AccessTokenURL: "https://login.live.com/oauth20_token.srf",
			AccessType:     "offline",
			UserinfoBase:   "https://apis.live.net",
			UserinfoURLs:   []string{"/v5.0/me"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].emails.account",
			PathName:       "url[0].name",
			PathPicture:    "",
			PathCover:      "",
		},
		LinkedIn: AuthRegistration{
			Network:        LinkedIn,
			Scopes:         []string{"r_basicprofile", "r_emailaddress"},
			AuthURL:        "https://www.linkedin.com/uas/oauth2/authorization",
			AccessTokenURL: "https://www.linkedin.com/uas/oauth2/accessToken",
			AccessType:     "offline",
			UserinfoURLs:   []string{"/v1/people/~:(picture-url,first-name,last-name,id,formatted-name,email-address)?format=json"},
			UserinfoBase:   "https://api.linkedin.com",
			PathID:         "url[0].id",
			PathEMail:      "url[0].emailAddress",
			PathName:       "url[0].formattedName",
			PathPicture:    "url[0].pictureUrl",
			PathCover:      "",
		},
	}
)

// GetRegistry returns a registry description for the given backend or an
// empty registration block.
func GetRegistry(backend Provider) AuthRegistration {
	res, _ := defaultBackends[backend]
	return res
}

// Instance returns a new registration provider with the given clientid
// and clientsecret.
func Instance(backend Provider, clientid, clientsecret string, scopes ...string) AuthRegistration {
	b := GetRegistry(backend)
	b.ClientID = clientid
	b.ClientSecret = clientsecret
	if scopes != nil && len(scopes) > 0 {
		b.Scopes = scopes
	}
	return b
}

// FillDefaults fills the given registration struct with the default values
// from the backend. The values are only overwritten if they are empty.
func FillDefaults(backend Provider, reg AuthRegistration) AuthRegistration {
	def := GetRegistry(backend)
	if reg.Scopes == nil || len(reg.Scopes) == 0 {
		reg.Scopes = def.Scopes
	}
	if reg.AuthURL == "" {
		reg.AuthURL = def.AuthURL
	}
	if reg.AccessTokenURL == "" {
		reg.AccessTokenURL = def.AccessTokenURL
	}
	if len(reg.UserinfoURLs) == 0 {
		reg.UserinfoURLs = def.UserinfoURLs
	}
	if reg.PathID == "" {
		reg.PathID = def.PathID
	}
	if reg.PathEMail == "" {
		reg.PathEMail = def.PathEMail
	}
	if reg.PathName == "" {
		reg.PathName = def.PathName
	}
	if reg.PathPicture == "" {
		reg.PathPicture = def.PathPicture
	}
	if reg.PathCover == "" {
		reg.PathCover = def.PathCover
	}
	return reg
}
<commit_msg>reduce the fields from facebook login<commit_after>package authkit

// some constants
const (
	Facebook Provider = "facebook"
	Google            = "google"
	Github            = "github"
	Live              = "live"
	LinkedIn          = "linkedin"
)

// A ProviderRegistry contains all registerd providers.
type ProviderRegistry map[Provider]AuthRegistration

var (
	defaultBackends = ProviderRegistry{
		Facebook: AuthRegistration{
			Network:        Facebook,
			Scopes:         []string{"public_profile", "email"},
			AuthURL:        "https://www.facebook.com/dialog/oauth/",
			AccessTokenURL: "https://graph.facebook.com/oauth/access_token",
			AccessType:     "offline",
			UserinfoBase:   "https://graph.facebook.com/v3.0",
			UserinfoURLs:   []string{"/me?fields=email,name,picture"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].email",
			PathName:       "url[0].name",
			PathPicture:    "url[0].picture.data.url",
			PathCover:      "",
		},
		Google: AuthRegistration{
			Network:        Google,
			Scopes:         []string{"openid", "profile", "email", "https://www.googleapis.com/auth/plus.me"},
			AuthURL:        "https://accounts.google.com/o/oauth2/auth",
			AccessTokenURL: "https://accounts.google.com/o/oauth2/token",
			AccessType:     "offline",
			UserinfoBase:   "https://www.googleapis.com",
			UserinfoURLs:   []string{"/plus/v1/people/me"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].emails[0].value",
			PathName:       "url[0].displayName",
			PathPicture:    "url[0].image.url",
			PathCover:      "url[0].cover.coverPhoto.url",
		},
		Github: AuthRegistration{
			Network:        Github,
			Scopes:         []string{"user:email"},
			AuthURL:        "https://github.com/login/oauth/authorize",
			AccessTokenURL: "https://github.com/login/oauth/access_token",
			AccessType:     "offline",
			UserinfoBase:   "https://api.github.com",
			UserinfoURLs:   []string{"/user", "/user/emails"},
			PathID:         "url[0].login",
			PathEMail:      "url[1].data[0].email",
			PathName:       "url[0].name",
			PathPicture:    "url[0].avatar_url",
			PathCover:      "",
		},
		Live: AuthRegistration{
			Network:        Live,
			Scopes:         []string{"wl.signin", "wl.emails", "wl.basic"},
			AuthURL:        "https://login.live.com/oauth20_authorize.srf",
			AccessTokenURL: "https://login.live.com/oauth20_token.srf",
			AccessType:     "offline",
			UserinfoBase:   "https://apis.live.net",
			UserinfoURLs:   []string{"/v5.0/me"},
			PathID:         "url[0].id",
			PathEMail:      "url[0].emails.account",
			PathName:       "url[0].name",
			PathPicture:    "",
			PathCover:      "",
		},
		LinkedIn: AuthRegistration{
			Network:        LinkedIn,
			Scopes:         []string{"r_basicprofile", "r_emailaddress"},
			AuthURL:        "https://www.linkedin.com/uas/oauth2/authorization",
			AccessTokenURL: "https://www.linkedin.com/uas/oauth2/accessToken",
			AccessType:     "offline",
			UserinfoURLs:   []string{"/v1/people/~:(picture-url,first-name,last-name,id,formatted-name,email-address)?format=json"},
			UserinfoBase:   "https://api.linkedin.com",
			PathID:         "url[0].id",
			PathEMail:      "url[0].emailAddress",
			PathName:       "url[0].formattedName",
			PathPicture:    "url[0].pictureUrl",
			PathCover:      "",
		},
	}
)

// GetRegistry returns a registry description for the given backend or an
// empty registration block.
func GetRegistry(backend Provider) AuthRegistration {
	res, _ := defaultBackends[backend]
	return res
}

// Instance returns a new registration provider with the given clientid
// and clientsecret.
func Instance(backend Provider, clientid, clientsecret string, scopes ...string) AuthRegistration {
	b := GetRegistry(backend)
	b.ClientID = clientid
	b.ClientSecret = clientsecret
	if scopes != nil && len(scopes) > 0 {
		b.Scopes = scopes
	}
	return b
}

// FillDefaults fills the given registration struct with the default values
// from the backend. The values are only overwritten if they are empty.
func FillDefaults(backend Provider, reg AuthRegistration) AuthRegistration {
	def := GetRegistry(backend)
	if reg.Scopes == nil || len(reg.Scopes) == 0 {
		reg.Scopes = def.Scopes
	}
	if reg.AuthURL == "" {
		reg.AuthURL = def.AuthURL
	}
	if reg.AccessTokenURL == "" {
		reg.AccessTokenURL = def.AccessTokenURL
	}
	if len(reg.UserinfoURLs) == 0 {
		reg.UserinfoURLs = def.UserinfoURLs
	}
	if reg.PathID == "" {
		reg.PathID = def.PathID
	}
	if reg.PathEMail == "" {
		reg.PathEMail = def.PathEMail
	}
	if reg.PathName == "" {
		reg.PathName = def.PathName
	}
	if reg.PathPicture == "" {
		reg.PathPicture = def.PathPicture
	}
	if reg.PathCover == "" {
		reg.PathCover = def.PathCover
	}
	return reg
}
<|endoftext|>


<commit_before>package victor

import (
    "github.com/brettbuddin/victor/utils/google"
    "strconv"
)

func registerDefaultAbilities(brain *Brain) {
    brain.Respond("campfire id", func(ctx *Context) {
        id := strconv.Itoa(ctx.Message().Id())
        ctx.Reply(id)
    })

    brain.Respond("ping", func(ctx *Context) {
        ctx.Reply("pong!")
    })

    brain.Respond("(image|img|gif|animate) (.*)", func(ctx *Context) {
        gifOnly := (ctx.Matches()[0] == "gif" || ctx.Matches()[0] == "animate")

        result, err := google.ImageSearch(ctx.Matches()[1], gifOnly)

        if err != nil {
            ctx.Send("There was error making the request.")
            return
        }

        if result == "" {
            ctx.Send("I didn't find anything.")
            return
        }

        ctx.Send(result)
    })

    brain.Respond("users (list|show)", func(ctx *Context) {
        list := ""

        for _, user := range brain.Users() {
            list += user.Name() + "\n"
        }

        ctx.Paste(list)
    })
}
<commit_msg>this was returning the message ID<commit_after>package victor

import (
    "github.com/brettbuddin/victor/utils/google"
    "strconv"
)

func registerDefaultAbilities(brain *Brain) {
    brain.Respond("me id", func(ctx *Context) {
        id := strconv.Itoa(ctx.Message().UserId())
        ctx.Reply(id)
    })

    brain.Respond("ping", func(ctx *Context) {
        ctx.Reply("pong!")
    })

    brain.Respond("(image|img|gif|animate) (.*)", func(ctx *Context) {
        gifOnly := (ctx.Matches()[0] == "gif" || ctx.Matches()[0] == "animate")

        result, err := google.ImageSearch(ctx.Matches()[1], gifOnly)

        if err != nil {
            ctx.Send("There was error making the request.")
            return
        }

        if result == "" {
            ctx.Send("I didn't find anything.")
            return
        }

        ctx.Send(result)
    })

    brain.Respond("users (list|show)", func(ctx *Context) {
        list := ""

        for _, user := range brain.Users() {
            list += user.Name() + "\n"
        }

        ctx.Paste(list)
    })
}
<|endoftext|>


<commit_before>package diff

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

	"github.com/omniscale/imposm3/cache"
	"github.com/omniscale/imposm3/config"
	"github.com/omniscale/imposm3/diff/download"
	"github.com/omniscale/imposm3/expire"
	"github.com/omniscale/imposm3/geom/limit"
	"github.com/omniscale/imposm3/logging"
)

var logger = logging.NewLogger("")

func Run() {
	if config.BaseOptions.Quiet {
		logging.SetQuiet(true)
	}

	var geometryLimiter *limit.Limiter
	if config.BaseOptions.LimitTo != "" {
		var err error
		step := logger.StartStep("Reading limitto geometries")
		geometryLimiter, err = limit.NewFromGeoJSON(
			config.BaseOptions.LimitTo,
			config.BaseOptions.LimitToCacheBuffer,
			config.BaseOptions.Srid,
		)
		if err != nil {
			logger.Fatal(err)
		}
		logger.StopStep(step)
	}
	downloader, err := download.NewDiffDownload(config.BaseOptions.DiffDir,
		config.BaseOptions.ReplicationUrl, config.BaseOptions.ReplicationInterval)
	if err != nil {
		logger.Fatal("unable to start diff downloader", err)
	}

	osmCache := cache.NewOSMCache(config.BaseOptions.CacheDir)
	err = osmCache.Open()
	if err != nil {
		logger.Fatal("osm cache: ", err)
	}
	defer osmCache.Close()

	diffCache := cache.NewDiffCache(config.BaseOptions.CacheDir)
	err = diffCache.Open()
	if err != nil {
		logger.Fatal("diff cache: ", err)
	}
	defer diffCache.Close()

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP)

	shutdown := func() {
		logger.Print("Exiting. (SIGTERM/SIGINT/SIGHUB)")
		logging.Shutdown()
		osmCache.Close()
		diffCache.Close()
		os.Exit(0)
	}

	var tiles *expire.TileList
	var tileExpireor expire.Expireor
	if config.BaseOptions.ExpireTilesDir != "" {
		tiles = expire.NewTileList(config.BaseOptions.ExpireTilesZoom, config.BaseOptions.ExpireTilesDir)
		tileExpireor = tiles
	}

	exp := newExpBackoff(2*time.Second, 5*time.Minute)

	for {
		select {
		case <-sigc:
			shutdown()
		case nextDiff := <-downloader.NextDiff:
			fname := nextDiff.FileName
			state := nextDiff.State
			for {
				p := logger.StartStep(fmt.Sprintf("importing #%d till %s", state.Sequence, state.Time))

				err := Update(fname, geometryLimiter, tileExpireor, osmCache, diffCache, false)

				osmCache.Coords.Flush()
				diffCache.Flush()

				if err == nil && tiles != nil {
					err := tiles.Flush()
					if err != nil {
						logger.Print("error writing tile expire list", err)
					}
				}

				logger.StopStep(p)

				select {
				case <-sigc:
					shutdown()
				default:
				}

				if err != nil {
					logger.Error(err)
					logger.Print("retrying in ", exp.Duration())
					exp.Wait()
				} else {
					exp.Reset()
					break
				}
			}
			if os.Getenv("IMPOSM3_SINGLE_DIFF") != "" {
				return
			}
		}
	}
}

type expBackoff struct {
	current time.Duration
	min     time.Duration
	max     time.Duration
}

func newExpBackoff(min, max time.Duration) *expBackoff {
	return &expBackoff{min, min, max}
}

func (eb *expBackoff) Duration() time.Duration {
	return eb.current
}

func (eb *expBackoff) Wait() {
	time.Sleep(eb.current)
	eb.current = eb.current * 2
	if eb.current > eb.max {
		eb.current = eb.max
	}
}

func (eb *expBackoff) Reset() {
	eb.current = eb.min
}
<commit_msg>flush tile list at most every 30 seconds<commit_after>package diff

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

	"github.com/omniscale/imposm3/cache"
	"github.com/omniscale/imposm3/config"
	"github.com/omniscale/imposm3/diff/download"
	"github.com/omniscale/imposm3/expire"
	"github.com/omniscale/imposm3/geom/limit"
	"github.com/omniscale/imposm3/logging"
)

var logger = logging.NewLogger("")

func Run() {
	if config.BaseOptions.Quiet {
		logging.SetQuiet(true)
	}

	var geometryLimiter *limit.Limiter
	if config.BaseOptions.LimitTo != "" {
		var err error
		step := logger.StartStep("Reading limitto geometries")
		geometryLimiter, err = limit.NewFromGeoJSON(
			config.BaseOptions.LimitTo,
			config.BaseOptions.LimitToCacheBuffer,
			config.BaseOptions.Srid,
		)
		if err != nil {
			logger.Fatal(err)
		}
		logger.StopStep(step)
	}
	downloader, err := download.NewDiffDownload(config.BaseOptions.DiffDir,
		config.BaseOptions.ReplicationUrl, config.BaseOptions.ReplicationInterval)
	if err != nil {
		logger.Fatal("unable to start diff downloader", err)
	}

	osmCache := cache.NewOSMCache(config.BaseOptions.CacheDir)
	err = osmCache.Open()
	if err != nil {
		logger.Fatal("osm cache: ", err)
	}
	defer osmCache.Close()

	diffCache := cache.NewDiffCache(config.BaseOptions.CacheDir)
	err = diffCache.Open()
	if err != nil {
		logger.Fatal("diff cache: ", err)
	}
	defer diffCache.Close()

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP)

	var tilelist *expire.TileList
	var lastTlFlush = time.Now()
	var tileExpireor expire.Expireor
	if config.BaseOptions.ExpireTilesDir != "" {
		tilelist = expire.NewTileList(config.BaseOptions.ExpireTilesZoom, config.BaseOptions.ExpireTilesDir)
		tileExpireor = tilelist
	}

	shutdown := func() {
		logger.Print("Exiting. (SIGTERM/SIGINT/SIGHUB)")
		logging.Shutdown()
		osmCache.Close()
		diffCache.Close()
		if tilelist != nil {
			err := tilelist.Flush()
			if err != nil {
				logger.Print("error writing tile expire list", err)
			}
		}
		os.Exit(0)
	}

	exp := newExpBackoff(2*time.Second, 5*time.Minute)

	for {
		select {
		case <-sigc:
			shutdown()
		case nextDiff := <-downloader.NextDiff:
			fname := nextDiff.FileName
			state := nextDiff.State
			for {
				p := logger.StartStep(fmt.Sprintf("importing #%d till %s", state.Sequence, state.Time))

				err := Update(fname, geometryLimiter, tileExpireor, osmCache, diffCache, false)

				osmCache.Coords.Flush()
				diffCache.Flush()

				if err == nil && tilelist != nil && time.Since(lastTlFlush) > time.Second*30 {
					// call at most once every 30 seconds to reduce files during the
					// catch-up phase after the initial import
					lastTlFlush = time.Now()
					err := tilelist.Flush()
					if err != nil {
						logger.Print("error writing tile expire list", err)
					}
				}

				logger.StopStep(p)

				select {
				case <-sigc:
					shutdown()
				default:
				}

				if err != nil {
					logger.Error(err)
					logger.Print("retrying in ", exp.Duration())
					exp.Wait()
				} else {
					exp.Reset()
					break
				}
			}
			if os.Getenv("IMPOSM3_SINGLE_DIFF") != "" {
				return
			}
		}
	}
}

type expBackoff struct {
	current time.Duration
	min     time.Duration
	max     time.Duration
}

func newExpBackoff(min, max time.Duration) *expBackoff {
	return &expBackoff{min, min, max}
}

func (eb *expBackoff) Duration() time.Duration {
	return eb.current
}

func (eb *expBackoff) Wait() {
	time.Sleep(eb.current)
	eb.current = eb.current * 2
	if eb.current > eb.max {
		eb.current = eb.max
	}
}

func (eb *expBackoff) Reset() {
	eb.current = eb.min
}
<|endoftext|>


<commit_before>package dskvs

import (
	"sync"
)

type page struct {
	isDirty   bool
	isDeleted bool
	basepath  string
	coll      string
	key       string
	value     []byte
	sync.RWMutex
}

func newPage(basepath, coll, key string) *page {
	return &page{
		isDirty:   false,
		isDeleted: false,
		basepath:  basepath,
		coll:      coll,
		key:       key,
		value:     nil,
	}
}

func (p *page) get() []byte {
	p.RLock()
	if p.isDeleted {
		return nil
	}
	p.RUnlock()
	return p.value
}

func (p *page) set(value []byte) {
	p.Lock()
	prev := len(p.value)
	next := len(value)
	if prev > next {
		p.value = p.value[:next]
	} else if prev < next {
		p.value = make([]byte, next)
	}
	copy(p.value, value)
	wasDirty := p.isDirty
	p.isDirty = true
	p.Unlock()
	if !wasDirty {
		jan.writePage(p)
	}
}

func (p *page) delete() {
	p.Lock()
	wasDirty := p.isDirty
	p.value = nil
	p.isDirty = true
	p.isDeleted = true
	p.Unlock()
	if !wasDirty {
		jan.writePage(p)
	}
}
<commit_msg>Safer `page.set` if client reads the bytes of `page.get` concurrently.<commit_after>package dskvs

import (
	"sync"
)

type page struct {
	isDirty   bool
	isDeleted bool
	basepath  string
	coll      string
	key       string
	value     []byte
	sync.RWMutex
}

func newPage(basepath, coll, key string) *page {
	return &page{
		isDirty:   false,
		isDeleted: false,
		basepath:  basepath,
		coll:      coll,
		key:       key,
		value:     nil,
	}
}

func (p *page) get() []byte {
	p.RLock()
	if p.isDeleted {
		return nil
	}
	p.RUnlock()
	return p.value
}

func (p *page) set(value []byte) {
	p.Lock()
	newBytes := make([]byte, len(value))
	copy(newBytes, value)
	p.value = newBytes
	wasDirty := p.isDirty
	p.isDirty = true
	p.Unlock()
	if !wasDirty {
		jan.writePage(p)
	}
}

func (p *page) delete() {
	p.Lock()
	wasDirty := p.isDirty
	p.value = nil
	p.isDirty = true
	p.isDeleted = true
	p.Unlock()
	if !wasDirty {
		jan.writePage(p)
	}
}
<|endoftext|>


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

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

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t * p, struct pcap_pkthdr **pkt_header,
		      u_char ** pkt_data)
{
	return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"

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

type Pcap struct {
	cptr *C.pcap_t
}

type pcapError struct{ string }

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

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

func Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

func Create(device string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

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

	h.cptr = C.pcap_create(dev, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}

	C.free(unsafe.Pointer(buf))
	return
}

// Set buffer size (units in bytes) on 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
}

//  If arg p is non-zero promiscuous mode will be set on capture handle when it is activated.
func (p *Pcap) SetPromisc(promisc bool) error {
	var pro int32
	if promisc {
		pro = 1
	}

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

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

// Set read timeout (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, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

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

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Openoffline
func OpenOffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

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

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var pkthdr C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec) * 1000) // pcap provides usec but time.Unix requires nsec
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

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

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

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

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

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

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

func FindAllDevs() (ifs []Interface, err string) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, C.GoString(buf)
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findAllAddresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}
		var a IFAddress
		var err error
		if a.IP, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.netmask))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddrToIP(rsa *syscall.RawSockaddr) (IP []byte, err 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
	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
	}
	err = errors.New("Unsupported address type")
	return
}

// Inject ...
func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

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

	if -1 == C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) {
		err = p.Geterror()
	}
	C.free(unsafe.Pointer(buf))
	return
}
<commit_msg>Added SetDirection function to allow clients to call the set_direction libpcap function<commit_after>// Package pcap is a wrapper around the pcap library.
package pcap

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

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t * p, struct pcap_pkthdr **pkt_header,
		      u_char ** pkt_data)
{
	return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"

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

type Pcap struct {
	cptr *C.pcap_t
}

type pcapError struct{ string }

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

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

func Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

func Create(device string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

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

	h.cptr = C.pcap_create(dev, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}

	C.free(unsafe.Pointer(buf))
	return
}

// Set buffer size (units in bytes) on 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
}

//  If arg p is non-zero promiscuous mode will be set on capture handle when it is activated.
func (p *Pcap) SetPromisc(promisc bool) error {
	var pro int32
	if promisc {
		pro = 1
	}

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

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

// Set read timeout (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, timeout_ms int32) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)
	var pro int32
	if promisc {
		pro = 1
	}

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

	h.cptr = C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Openoffline
func OpenOffline(file string) (handle *Pcap, err error) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	h := new(Pcap)

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

	h.cptr = C.pcap_open_offline(cf, buf)
	if nil == h.cptr {
		handle = nil
		err = &pcapError{C.GoString(buf)}
	} else {
		handle = h
	}
	C.free(unsafe.Pointer(buf))
	return
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (pkt *Packet, result int32) {
	var pkthdr_ptr *C.struct_pcap_pkthdr
	var pkthdr C.struct_pcap_pkthdr

	var buf_ptr *C.u_char
	var buf unsafe.Pointer
	result = int32(C.hack_pcap_next_ex(p.cptr, &pkthdr_ptr, &buf_ptr))

	buf = unsafe.Pointer(buf_ptr)
	pkthdr = *pkthdr_ptr

	if nil == buf {
		return
	}
	pkt = new(Packet)
	pkt.Time = time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec) * 1000) // pcap provides usec but time.Unix requires nsec
	pkt.Caplen = uint32(pkthdr.caplen)
	pkt.Len = uint32(pkthdr.len)
	pkt.Data = make([]byte, pkthdr.caplen)

	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return
}

func (p *Pcap) Getstats() (stat *Stat, err error) {
	var cstats _Ctype_struct_pcap_stat
	if -1 == C.pcap_stats(p.cptr, &cstats) {
		return nil, p.Geterror()
	}
	stats := new(Stat)
	stats.PacketsReceived = uint32(cstats.ps_recv)
	stats.PacketsDropped = uint32(cstats.ps_drop)
	stats.PacketsIfDropped = uint32(cstats.ps_ifdrop)

	return stats, nil
}

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

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

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

func (p *Pcap) SetDirection(direction string) (err error) {
    var pcap_direction C.pcap_direction_t
    if (direction == "in") {
        pcap_direction = C.PCAP_D_IN
    } else if (direction == "out") {
        pcap_direction = C.PCAP_D_OUT
    } else {
        pcap_direction = C.PCAP_D_INOUT
    }
    if -1 == C.pcap_setdirection(p.cptr, pcap_direction) {
        return p.Geterror()
    }
    return nil
}

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

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

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

func FindAllDevs() (ifs []Interface, err string) {
	var buf *C.char
	buf = (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))
	var alldevsp *C.pcap_if_t

	if -1 == C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) {
		return nil, C.GoString(buf)
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))
	dev := alldevsp
	var i uint32
	for i = 0; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		i++
	}
	ifs = make([]Interface, i)
	dev = alldevsp
	for j := uint32(0); dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		var iface Interface
		iface.Name = C.GoString(dev.name)
		iface.Description = C.GoString(dev.description)
		iface.Addresses = findAllAddresses(dev.addresses)
		// TODO: add more elements
		ifs[j] = iface
		j++
	}
	return
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) (retval []IFAddress) {
	// TODO - make it support more than IPv4 and IPv6?
	retval = make([]IFAddress, 0, 1)
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}
		var a IFAddress
		var err error
		if a.IP, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr))); err != nil {
			continue
		}
		if a.Netmask, err = sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.netmask))); err != nil {
			continue
		}
		retval = append(retval, a)
	}
	return
}

func sockaddrToIP(rsa *syscall.RawSockaddr) (IP []byte, err 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
	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
	}
	err = errors.New("Unsupported address type")
	return
}

// Inject ...
func (p *Pcap) Inject(data []byte) (err error) {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))

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

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


<commit_before>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	PieceBlock
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool

	requestPieceChannel   chan uint32
	sendPieceBlockChannel chan BlockMessage
	sendHaveChannel       chan uint32
	done                  chan struct{}

	id string
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

type BlockMessage struct {
	index uint32
	begin uint32
	block []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localChoked = true
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	addr := fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)

	var err error
	peer.connection, err = net.Dial("tcp4", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[0:20], []byte("\x13BitTorrent protocol")) {
		log.Printf("bad protocol from peer: %s\n", addr)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s\n", addr)
		return
	} else if len(peer.id) != 0 && !bytes.Equal(handshake[48:68], []byte(peer.id)) {
		log.Printf("peer id mismatch from peer: %s\n", addr)
		return
	}

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(&blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		peer.sendUnchokeMessage()
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		if !peer.remoteInterested {
			return errors.New("peer sent request without showing interest")
		}

		if peer.localChoked {
			return errors.New("peer sent request while choked")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		peer.torrent.blockRequestChannel <- BlockRequestMessage{peer, index, begin, length}
	case PieceBlock:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	if piece, _ := peer.getPeerPiece(index); piece != nil {
		return
	}

	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(9+len(blockMessage.block)))
	packet.WriteByte(PieceBlock)
	binary.Write(&packet, binary.BigEndian, blockMessage.index)
	binary.Write(&packet, binary.BigEndian, blockMessage.begin)
	packet.Write(blockMessage.block)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := bytes.NewBuffer([]byte{0, 0, 0, 5, Have})
	binary.Write(packet, binary.BigEndian, pieceIndex)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendUnchokeMessage() {
	if !peer.localChoked {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Unchoke})
	peer.localChoked = false
}
<commit_msg>Allow connecting to IPv6 addresses<commit_after>/*
 * Copyright (c) 2014 Mark Samman <https://github.com/marksamman/gotorrent>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"net"
)

const (
	Choke = iota
	Unchoke
	Interested
	Uninterested
	Have
	Bitfield
	Request
	PieceBlock
	Cancel
	Port
)

type Peer struct {
	ip      net.IP
	port    uint16
	torrent *Torrent

	pieces           []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool

	requestPieceChannel   chan uint32
	sendPieceBlockChannel chan BlockMessage
	sendHaveChannel       chan uint32
	done                  chan struct{}

	id string
}

type PeerPiece struct {
	index     uint32
	data      []byte
	writes    int
	reqWrites int
}

type Packet struct {
	length      uint32
	messageType int
	payload     []byte
}

type BlockMessage struct {
	index uint32
	begin uint32
	block []byte
}

func NewPeer(torrent *Torrent) *Peer {
	peer := Peer{}
	peer.torrent = torrent

	peer.remoteChoked = true
	peer.remoteInterested = false
	peer.localChoked = true
	peer.localInterested = false
	return &peer
}

func (peer *Peer) readN(n int) ([]byte, error) {
	buf := make([]byte, n)
	for pos := 0; pos < n; {
		count, err := peer.connection.Read(buf[pos:])
		if err != nil {
			return nil, err
		}
		pos += count
	}
	return buf, nil
}

func (peer *Peer) connect() {
	var addr string
	if peer.ip.To4() != nil {
		addr = fmt.Sprintf("%s:%d", peer.ip.String(), peer.port)
	} else {
		addr = fmt.Sprintf("[%s]:%d", peer.ip.String(), peer.port)
	}

	var err error
	peer.connection, err = net.Dial("tcp", addr)
	if err != nil {
		log.Printf("failed to connect to peer: %s\n", err)
		return
	}
	defer peer.connection.Close()

	log.Printf("connected to peer: %s\n", addr)

	// Send handshake
	if _, err := peer.connection.Write(peer.torrent.handshake); err != nil {
		log.Printf("failed to send handshake to peer: %s\n", err)
		return
	}

	// Receive handshake
	if handshake, err := peer.readN(68); err != nil {
		log.Printf("failed to read handshake from peer: %s\n", err)
		return
	} else if !bytes.Equal(handshake[0:20], []byte("\x13BitTorrent protocol")) {
		log.Printf("bad protocol from peer: %s\n", addr)
		return
	} else if !bytes.Equal(handshake[28:48], peer.torrent.infoHash) {
		log.Printf("info hash mismatch from peer: %s\n", addr)
		return
	} else if len(peer.id) != 0 && !bytes.Equal(handshake[48:68], []byte(peer.id)) {
		log.Printf("peer id mismatch from peer: %s\n", addr)
		return
	}

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})

	peer.torrent.addPeerChannel <- peer
	defer func() {
		peer.torrent.removePeerChannel <- peer
	}()

	packetChannel := make(chan Packet)
	errorChannel := make(chan error)

	go peer.receiver(packetChannel, errorChannel)
	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(&blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case packet := <-packetChannel:
			if err := peer.processMessage(&packet); err != nil {
				log.Printf("error while processing message in peer %s: %s", addr, err)
				return
			}
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			return
		case <-peer.done:
			return
		}
	}
}

func (peer *Peer) receiver(packetChannel chan Packet, errorChannel chan error) {
	for {
		lengthHeader, err := peer.readN(4)
		if err != nil {
			errorChannel <- err
			return
		}

		var length uint32
		binary.Read(bytes.NewBuffer(lengthHeader), binary.BigEndian, &length)
		if length == 0 {
			// keep-alive
			continue
		}

		data, err := peer.readN(int(length))
		if err != nil {
			errorChannel <- err
			continue
		}

		packetChannel <- Packet{length, int(data[0]), data[1:]}
	}
}

func (peer *Peer) processMessage(packet *Packet) error {
	switch packet.messageType {
	case Choke:
		if packet.length != 1 {
			return errors.New("length of choke packet must be 1")
		}

		peer.remoteChoked = true
	case Unchoke:
		if packet.length != 1 {
			return errors.New("length of unchoke packet must be 1")
		}

		peer.remoteChoked = false
		for _, piece := range peer.pieces {
			peer.requestPiece(piece)
		}
	case Interested:
		if packet.length != 1 {
			return errors.New("length of interested packet must be 1")
		}

		peer.remoteInterested = true
		peer.sendUnchokeMessage()
	case Uninterested:
		if packet.length != 1 {
			return errors.New("length of not interested packet must be 1")
		}

		peer.remoteInterested = false
	case Have:
		if packet.length != 5 {
			return errors.New("length of have packet must be 5")
		}

		var index uint32
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &index)
		peer.torrent.havePieceChannel <- HavePieceMessage{peer, index}
	case Bitfield:
		if packet.length < 2 {
			return errors.New("length of bitfield packet must be at least 2")
		}
		peer.torrent.bitfieldChannel <- BitfieldMessage{peer, packet.payload}
	case Request:
		if packet.length != 13 {
			return errors.New("length of request packet must be 13")
		}

		if !peer.remoteInterested {
			return errors.New("peer sent request without showing interest")
		}

		if peer.localChoked {
			return errors.New("peer sent request while choked")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		if length > 32768 {
			return errors.New("peer requested length over 32KB")
		}
		peer.torrent.blockRequestChannel <- BlockRequestMessage{peer, index, begin, length}
	case PieceBlock:
		if packet.length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		var index, begin uint32

		buf := bytes.NewBuffer(packet.payload)
		binary.Read(buf, binary.BigEndian, &index)

		piece, idx := peer.getPeerPiece(index)
		if piece == nil {
			return errors.New("received index we didn't ask for")
		}

		binary.Read(buf, binary.BigEndian, &begin)

		if int64(begin)+int64(packet.length)-9 > int64(len(piece.data)) {
			return errors.New("begin+length exceeds length of data buffer")
		}

		copy(piece.data[begin:], packet.payload[8:])
		piece.writes++

		if piece.writes == piece.reqWrites {
			// Send piece to Torrent
			peer.torrent.pieceChannel <- PieceMessage{peer, index, piece.data}

			// Remove piece from peer
			peer.pieces = append(peer.pieces[:idx], peer.pieces[idx+1:]...)
		}
	case Cancel:
		if packet.length != 13 {
			return errors.New("length of cancel packet must be 13")
		}

		buf := bytes.NewBuffer(packet.payload)
		var index, begin, length uint32
		binary.Read(buf, binary.BigEndian, &index)
		binary.Read(buf, binary.BigEndian, &begin)
		binary.Read(buf, binary.BigEndian, &length)
		// TODO: Handle cancel
	case Port:
		if packet.length != 3 {
			return errors.New("length of port packet must be 3")
		}

		var port uint16
		binary.Read(bytes.NewBuffer(packet.payload), binary.BigEndian, &port)
		// Peer has a DHT node on port
	}
	return nil
}

func (peer *Peer) sendInterested() {
	if peer.localInterested {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Interested})
	peer.localInterested = true
}

func (peer *Peer) sendRequest(index, begin, length uint32) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(13))
	packet.WriteByte(Request)
	binary.Write(&packet, binary.BigEndian, index)
	binary.Write(&packet, binary.BigEndian, begin)
	binary.Write(&packet, binary.BigEndian, length)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) getPeerPiece(index uint32) (*PeerPiece, int) {
	for idx, piece := range peer.pieces {
		if piece.index == index {
			return piece, idx
		}
	}
	return nil, -1
}

func (peer *Peer) sendPieceRequest(index uint32) {
	if piece, _ := peer.getPeerPiece(index); piece != nil {
		return
	}

	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(int(index))
	peer.pieces = append(peer.pieces, &PeerPiece{index, make([]byte, pieceLength), 0, int(math.Ceil(float64(pieceLength) / 16384))})
	if !peer.remoteChoked {
		peer.requestPiece(peer.pieces[len(peer.pieces)-1])
	}
}

func (peer *Peer) requestPiece(piece *PeerPiece) {
	var pos uint32
	pieceLength := uint32(len(piece.data))
	for pieceLength > 16384 {
		peer.sendRequest(piece.index, pos, 16384)
		pieceLength -= 16384
		pos += 16384
	}
	peer.sendRequest(piece.index, pos, pieceLength)
}

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	var packet bytes.Buffer
	binary.Write(&packet, binary.BigEndian, uint32(9+len(blockMessage.block)))
	packet.WriteByte(PieceBlock)
	binary.Write(&packet, binary.BigEndian, blockMessage.index)
	binary.Write(&packet, binary.BigEndian, blockMessage.begin)
	packet.Write(blockMessage.block)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := bytes.NewBuffer([]byte{0, 0, 0, 5, Have})
	binary.Write(packet, binary.BigEndian, pieceIndex)
	peer.connection.Write(packet.Bytes())
}

func (peer *Peer) sendUnchokeMessage() {
	if !peer.localChoked {
		return
	}

	peer.connection.Write([]byte{0, 0, 0, 1, Unchoke})
	peer.localChoked = false
}
<|endoftext|>


<commit_before>package quibit

import (
	"fmt"
	"net"
	"strconv"
	"time"
)

type Peer struct {
	IP   net.IP
	Port uint16
	conn net.Conn
}

func peerFromConn(conn net.Conn) Peer {
	var p Peer
	addr := conn.RemoteAddr()
	if addr.Network() != "tcp" {
		return p
	}
	ip, portStr, err := net.SplitHostPort(addr.String())
	port, _ := strconv.Atoi(portStr)
	if err != nil {
		return p
	}

	// Create New Peer
	p.IP = net.ParseIP(ip)
	p.Port = uint16(port)

	p.conn = conn
	return p
}

func (p *Peer) String() string {
	if p == nil {
		return ""
	}
	return net.JoinHostPort(p.IP.String(), strconv.Itoa(int(p.Port)))
}

func (p *Peer) IsConnected() bool {
	if p == nil {
		return false
	}
	return (p.conn == nil)
}

func (p *Peer) connect() error {
	// Check for sane peer object
	if p == nil {
		return QuibitError(eNILOBJ)
	}
	if p.conn != nil {
		return nil
	}

	var err error
	p.conn, err = net.DialTimeout("tcp", p.String(), 10 * time.Second)
	if err != nil {
		p.conn = nil
		return err
	}

	// Set Keep-Alives
	p.conn.(*net.TCPConn).SetKeepAlive(true)
	p.conn.(*net.TCPConn).SetKeepAlivePeriod(time.Second)

	return nil
}

func (p *Peer) Disconnect() {
	if p.conn == nil {
		fmt.Println("Peer already disconnected!")
		return
	}
	p.conn.Close()
	p.conn = nil
}

func (p *Peer) sendFrame(frame Frame) error {
	var n int
	if p == nil {
		return QuibitError(eNILOBJ)
	}
	if p.conn == nil {
		return QuibitError(eNILOBJ)
	}

	var err error

	headerBytes, err := frame.Header.ToBytes()
	if err != nil {
		return QuibitError(eHEADER)
	}

	n, err = p.conn.Write(append(headerBytes, frame.Payload...))
	fmt.Printf("Wrote %d Bytes to Peer: %s\n", n, p)
	if err != nil {
		return err
	}

	return nil
}

func (p *Peer) receive(recvChan chan Frame, log chan string) {
	if p.conn == nil {
		return
	}
	for {
		// So now we have a connection.  Let's shake hands.
		frame, err := recvAll(p.conn, log)
		if err != nil {
			fmt.Println("Error receiving from Peer: ", err)
			p.Disconnect()
			break
		} else {
			frame.Peer = p.String()
			fmt.Println("Sending to Recv Channel from... ", p)
			recvChan <- frame
		}
	} // End for
} // End receive()
<commit_msg>Changed debugging<commit_after>package quibit

import (
	"fmt"
	"net"
	"strconv"
	"time"
)

type Peer struct {
	IP   net.IP
	Port uint16
	conn net.Conn
}

func peerFromConn(conn net.Conn) Peer {
	var p Peer
	addr := conn.RemoteAddr()
	if addr.Network() != "tcp" {
		return p
	}
	ip, portStr, err := net.SplitHostPort(addr.String())
	port, _ := strconv.Atoi(portStr)
	if err != nil {
		return p
	}

	// Create New Peer
	p.IP = net.ParseIP(ip)
	p.Port = uint16(port)

	p.conn = conn
	return p
}

func (p *Peer) String() string {
	if p == nil {
		return ""
	}
	return net.JoinHostPort(p.IP.String(), strconv.Itoa(int(p.Port)))
}

func (p *Peer) IsConnected() bool {
	if p == nil {
		return false
	}
	return (p.conn == nil)
}

func (p *Peer) connect() error {
	// Check for sane peer object
	if p == nil {
		return QuibitError(eNILOBJ)
	}
	if p.conn != nil {
		return nil
	}

	var err error
	p.conn, err = net.DialTimeout("tcp", p.String(), 10 * time.Second)
	if err != nil {
		p.conn = nil
		return err
	}

	// Set Keep-Alives
	p.conn.(*net.TCPConn).SetKeepAlive(true)
	p.conn.(*net.TCPConn).SetKeepAlivePeriod(time.Second)

	return nil
}

func (p *Peer) Disconnect() {
	if p.conn == nil {
		fmt.Println("Peer already disconnected!")
		return
	}
	p.conn.Close()
	p.conn = nil
}

func (p *Peer) sendFrame(frame Frame) error {
	var n int
	if p == nil {
		return QuibitError(eNILOBJ)
	}
	if p.conn == nil {
		return QuibitError(eNILOBJ)
	}

	var err error

	headerBytes, err := frame.Header.ToBytes()
	if err != nil {
		return QuibitError(eHEADER)
	}

	n, err = p.conn.Write(append(headerBytes, frame.Payload...))
	fmt.Printf("Wrote %d Bytes to Peer: %s\n", n, p)
	if err != nil {
		return err
	}

	return nil
}

func (p *Peer) receive(recvChan chan Frame, log chan string) {
	if p.conn == nil {
		return
	}
	for {
		// So now we have a connection.  Let's shake hands.
		frame, err := recvAll(p.conn, log)
		if err != nil {
			fmt.Println("Error receiving from Peer: ", err)
			p.Disconnect()
			break
		} else {
			frame.Peer = p.String()
			fmt.Println("Sending to Recv Channel from... ", p.conn.RemoteAddr())
			recvChan <- frame
		}
	} // End for
} // End receive()
<|endoftext|>


<commit_before>package pick

import (
	"io"
	"strings"

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

type PickOption struct {
	PageSource string
	TagName    string
	Attr       *Attr //optional
}

type Attr struct {
	Label string
	Value string
}

func PickAttr(option *PickOption, AttrLabel string) (data []string, err error) {
	if option == nil || option.PageSource == "" {
		return data, nil
	}

	z := html.NewTokenizer(strings.NewReader(option.PageSource))

	for {
		tokenType := z.Next()

		switch tokenType {

		//ignore the error token
		//quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return data, nil
			}

		case html.SelfClosingTagToken:
			fallthrough
		case html.StartTagToken:
			tagName, attr := z.TagName()

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			attr_arr := []*Attr{}

			matched := false

			//get attr
			for attr {
				label, value, attr = z.TagAttr()

				label_str := string(label)
				value_str := string(value)

				if option.Attr == nil || (option.Attr.Label == label_str && option.Attr.Value == value_str) {
					matched = true
				}

				attr_arr = append(attr_arr, &Attr{
					label_str,
					value_str,
				})
			}

			if !matched {
				continue
			}

			//loop attr
			for i := 0; i < len(attr_arr); i++ {
				attr := attr_arr[i]

				if attr.Label == AttrLabel {
					data = append(data, attr.Value)
				}
			}
		}
	}

	return data, z.Err()
}
<commit_msg>refactor, reduce 1 loop<commit_after>package pick

import (
	"io"
	"strings"

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

type PickOption struct {
	PageSource string
	TagName    string
	Attr       *Attr //optional
}

type Attr struct {
	Label string
	Value string
}

func PickAttr(option *PickOption, AttrLabel string) (data []string, err error) {
	if option == nil || option.PageSource == "" {
		return data, nil
	}

	z := html.NewTokenizer(strings.NewReader(option.PageSource))

	for {
		tokenType := z.Next()

		switch tokenType {

		//ignore the error token
		//quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return data, nil
			}

		case html.SelfClosingTagToken:
			fallthrough
		case html.StartTagToken:
			tagName, attr := z.TagName()

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			data_tmp := []string{}

			matched := false

			//get attr
			for attr {
				label, value, attr = z.TagAttr()

				label_str := string(label)
				value_str := string(value)

				if option.Attr == nil || (option.Attr.Label == label_str && option.Attr.Value == value_str) {
					matched = true
				}

				if label_str == AttrLabel {
					data_tmp = append(data_tmp, value_str)
				}
			}

			if !matched {
				continue
			}

			//loop attr
			data = append(data, data_tmp...)
		}
	}

	return data, z.Err()
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Rene Richard

This file is part of zmq-soundtouch.

zmq-soundtouch is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

zmq-soundtouch is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with zmq-soundtouch.  If not, see <http://www.gnu.org/licenses/>.
*/

package messaging

import (
	//	"errors"
	"fmt"
	zmq "github.com/pebbe/zmq4"
	"strings"
)

type dealer struct {
	ctx      *zmq.Context
	msgChan  chan string
	doneChan chan bool
	errChan  chan error
	client   *zmq.Socket
}

func NewDealer() *dealer {
	ctx, _ := zmq.NewContext()

	msgChan := make(chan string)
	doneChan := make(chan bool)
	errChan := make(chan error)

	client, err := ctx.NewSocket(zmq.DEALER)

	//TODO fix...
	if err != nil {
		fmt.Println(err)
	}

	return &dealer{ctx, msgChan, doneChan, errChan, client}
}

func (d *dealer) GetCacheContent() {
	d.msgChan <- "ICANHAZ"
}

func (d *dealer) Error(err error) {
	d.errChan <- err
}

func (d *dealer) Done() {
	d.doneChan <- true
}

func (d *dealer) readMessages() {

	for {
		select {

		//We receive a message on the message channel
		case msg := <-d.msgChan:
			fmt.Println("received : ", msg)

			if strings.Compare(msg, "KTHXBYE") == 0 {
				d.doneChan <- true // for recieveMessages method
				return
			} else {
				fmt.Println("Got Message")
				fmt.Println(msg)
			}
		//We have an error
		case err := <-d.errChan:
			d.errChan <- err // for receiveMessages method
			return
		//We're done
		case <-d.doneChan:
			d.doneChan <- true // for receiveMessages method
			return
		}

	}
}

func (d *dealer) receiveMessages() {

	for {
		select {
		//We have an error
		case err := <-d.errChan:
			d.errChan <- err // for readMessages method
			return

		//We're done
		case <-d.doneChan:
			d.doneChan <- true // for readMessages method
			return

		// read data from socket connection (loop)
		default:
			reply, err := d.client.Recv(0)

			if err != nil {
				d.errChan <- err
			}

			d.msgChan <- reply
		}
	}

}

func (d *dealer) Start() {

	defer d.ctx.Term()

	//TODO log
	fmt.Println("Starting...")

	d.client.Connect("tcp://127.0.0.1:8000")
	defer d.client.Close()

	d.client.Send("ICANHAZ?", 0)

	go d.readMessages()
	d.receiveMessages()

}

/*
func (d *dealer) Start() {
	defer d.ctx.Term()

	//TODO log
	fmt.Println("Starting...")
	client, err := d.ctx.NewSocket(zmq.DEALER)

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

	client.Connect("tcp://127.0.0.1:8000")
	defer client.Close()

	client.Send("ICANHAZ?", 0)

	for {
		reply, _ := client.Recv(0)
		fmt.Println("received : ", reply)
		if strings.Compare(reply, "KTHXBYE") == 0 {
			return
		}
	}

}*/
<commit_msg>zmqdealer updates<commit_after>/*
Copyright 2016 Rene Richard

This file is part of zmq-soundtouch.

zmq-soundtouch is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

zmq-soundtouch is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with zmq-soundtouch.  If not, see <http://www.gnu.org/licenses/>.
*/

package messaging

import (
	//	"errors"
	"fmt"
	zmq "github.com/pebbe/zmq4"
	"strings"
)

type dealer struct {
	ctx      *zmq.Context
	msgChan  chan string
	doneChan chan bool
	errChan  chan error
	client   *zmq.Socket
}

func NewDealer() *dealer {
	ctx, _ := zmq.NewContext()

	msgChan := make(chan string)
	doneChan := make(chan bool)
	errChan := make(chan error)

	client, err := ctx.NewSocket(zmq.DEALER)

	//TODO fix...
	if err != nil {
		fmt.Println(err)
	}

	return &dealer{ctx, msgChan, doneChan, errChan, client}
}

func (d *dealer) GetCacheContent() {
	d.msgChan <- "ICANHAZ"
}

func (d *dealer) Error(err error) {
	d.errChan <- err
}

func (d *dealer) Done() {
	d.doneChan <- true
}

func (d *dealer) readMessages() {

	for {
		select {

		//We receive a message on the message channel
		case msg := <-d.msgChan:
			fmt.Println("received : ", msg)

			if strings.Compare(msg, "KTHXBYE") == 0 {
				d.doneChan <- true // for recieveMessages method
				return
			} else {
				fmt.Println("Got Message")
				fmt.Println(msg)
				return
			}
		//We have an error
		case err := <-d.errChan:
			d.errChan <- err // for receiveMessages method
			return
		//We're done
		case <-d.doneChan:
			d.doneChan <- true // for receiveMessages method
			return
		}

	}
}

func (d *dealer) receiveMessages() {

	for {
		select {
		//We have an error
		case err := <-d.errChan:
			d.errChan <- err // for readMessages method
			return

		//We're done
		case <-d.doneChan:
			d.doneChan <- true // for readMessages method
			return

		// read data from socket connection (loop)
		default:
			reply, err := d.client.Recv(0)

			if err != nil {
				d.errChan <- err
			}

			d.msgChan <- reply
		}
	}

}

func (d *dealer) Start() {

	defer d.ctx.Term()

	//TODO log
	fmt.Println("Starting...")

	d.client.Connect("tcp://127.0.0.1:8000")
	defer d.client.Close()

	d.client.Send("ICANHAZ?", 0)

	go d.readMessages()
	d.receiveMessages()

}

/*
func (d *dealer) Start() {
	defer d.ctx.Term()

	//TODO log
	fmt.Println("Starting...")
	client, err := d.ctx.NewSocket(zmq.DEALER)

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

	client.Connect("tcp://127.0.0.1:8000")
	defer client.Close()

	client.Send("ICANHAZ?", 0)

	for {
		reply, _ := client.Recv(0)
		fmt.Println("received : ", reply)
		if strings.Compare(reply, "KTHXBYE") == 0 {
			return
		}
	}

}*/
<|endoftext|>


<commit_before>package hypervisor

import (
	"net"
	"time"

	"github.com/Symantec/Dominator/lib/tags"
)

const (
	StateStarting      = 0
	StateRunning       = 1
	StateFailedToStart = 2
	StateStopping      = 3
	StateStopped       = 4
	StateDestroying    = 5
	StateMigrating     = 6

	VolumeFormatRaw   = 0
	VolumeFormatQCOW2 = 1
)

type AcknowledgeVmRequest struct {
	IpAddress net.IP
}

type AcknowledgeVmResponse struct {
	Error string
}

type Address struct {
	IpAddress  net.IP `json:",omitempty"`
	MacAddress string
}

type BecomePrimaryVmOwnerRequest struct {
	IpAddress net.IP
}

type BecomePrimaryVmOwnerResponse struct {
	Error string
}

type ChangeAddressPoolRequest struct {
	AddressesToAdd       []Address
	MaximumFreeAddresses map[string]uint // Key: subnet ID.
}

type ChangeAddressPoolResponse struct {
	Error string
}

type ChangeOwnersRequest struct {
	OwnerGroups []string `json:",omitempty"`
	OwnerUsers  []string `json:",omitempty"`
}

type ChangeOwnersResponse struct {
	Error string
}

type ChangeVmOwnerUsersRequest struct {
	IpAddress  net.IP
	OwnerUsers []string
}

type ChangeVmOwnerUsersResponse struct {
	Error string
}

type ChangeVmTagsRequest struct {
	IpAddress net.IP
	Tags      tags.Tags
}

type ChangeVmTagsResponse struct {
	Error string
}

type CommitImportedVmRequest struct {
	IpAddress net.IP
}

type CommitImportedVmResponse struct {
	Error string
}

type CreateVmRequest struct {
	DhcpTimeout      time.Duration
	ImageDataSize    uint64
	ImageTimeout     time.Duration
	MinimumFreeBytes uint64
	RoundupPower     uint64
	SecondaryVolumes []Volume
	UserDataSize     uint64
	VmInfo
} // RAW image data (length=ImageDataSize) and user data (length=UserDataSize)
// are streamed afterwards.

type CreateVmResponse struct { // Multiple responses are sent.
	DhcpTimedOut    bool
	Final           bool // If true, this is the final response.
	IpAddress       net.IP
	ProgressMessage string
	Error           string
}

type DeleteVmVolumeRequest struct {
	AccessToken []byte
	IpAddress   net.IP
	VolumeIndex uint
}

type DeleteVmVolumeResponse struct {
	Error string
}

type DestroyVmRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type DestroyVmResponse struct {
	Error string
}

type DiscardVmAccessTokenRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type DiscardVmAccessTokenResponse struct {
	Error string
}

type DiscardVmOldImageRequest struct {
	IpAddress net.IP
}

type DiscardVmOldImageResponse struct {
	Error string
}

type DiscardVmOldUserDataRequest struct {
	IpAddress net.IP
}

type DiscardVmOldUserDataResponse struct {
	Error string
}

type DiscardVmSnapshotRequest struct {
	IpAddress net.IP
}

type DiscardVmSnapshotResponse struct {
	Error string
}

// The GetUpdates() RPC is fully streamed.
// The client may or may not send GetUpdateRequest messages to the server.
// The server sends a stream of Update messages.

type GetUpdateRequest struct{}

type Update struct {
	HaveAddressPool  bool               `json:",omitempty"`
	AddressPool      []Address          `json:",omitempty"` // Used & free.
	NumFreeAddresses map[string]uint    `json:",omitempty"` // Key: subnet ID.
	HealthStatus     string             `json:",omitempty"`
	HaveSerialNumber bool               `json:",omitempty"`
	SerialNumber     string             `json:",omitempty"`
	HaveSubnets      bool               `json:",omitempty"`
	Subnets          []Subnet           `json:",omitempty"`
	HaveVMs          bool               `json:",omitempty"`
	VMs              map[string]*VmInfo `json:",omitempty"` // Key: IP address.
}

type GetVmAccessTokenRequest struct {
	IpAddress net.IP
	Lifetime  time.Duration
}

type GetVmAccessTokenResponse struct {
	Token []byte `json:",omitempty"`
	Error string
}

type GetVmInfoRequest struct {
	IpAddress net.IP
}

type GetVmInfoResponse struct {
	VmInfo VmInfo
	Error  string
}

type GetVmUserDataRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type GetVmUserDataResponse struct {
	Error  string
	Length uint64
} // Data (length=Length) are streamed afterwards.

// The GetVmVolume() RPC is followed by the proto/rsync.GetBlocks message.

type GetVmVolumeRequest struct {
	AccessToken []byte
	IpAddress   net.IP
	VolumeIndex uint
}

type GetVmVolumeResponse struct {
	Error string
}

type ImportLocalVmRequest struct {
	VerificationCookie []byte `json:",omitempty"`
	VmInfo
	VolumeFilenames []string
}

type ImportLocalVmResponse struct {
	Error string
}

type ListVMsRequest struct {
	Sort bool
}

type ListVMsResponse struct {
	IpAddresses []net.IP
}

type ListVolumeDirectoriesRequest struct{}

type ListVolumeDirectoriesResponse struct {
	Directories []string
	Error       string
}

type MigrateVmRequest struct {
	AccessToken      []byte
	DhcpTimeout      time.Duration
	IpAddress        net.IP
	SourceHypervisor string
}

type MigrateVmResponse struct { // Multiple responses are sent.
	Error           string
	Final           bool // If true, this is the final response.
	ProgressMessage string
	RequestCommit   bool
}

type MigrateVmResponseResponse struct {
	Commit bool
}

type NetbootMachineRequest struct {
	Address                      Address
	Files                        map[string][]byte
	FilesExpiration              time.Duration
	Hostname                     string
	NumAcknowledgementsToWaitFor uint
	OfferExpiration              time.Duration
	Subnet                       *Subnet
	WaitTimeout                  time.Duration
}

type NetbootMachineResponse struct {
	Error string
}

type PrepareVmForMigrationRequest struct {
	AccessToken []byte
	Enable      bool
	IpAddress   net.IP
}

type PrepareVmForMigrationResponse struct {
	Error string
}

type ProbeVmPortRequest struct {
	IpAddress  net.IP
	PortNumber uint
	Timeout    time.Duration
}

type ProbeVmPortResponse struct {
	PortIsOpen bool
	Error      string
}

type ReplaceVmImageRequest struct {
	DhcpTimeout      time.Duration
	ImageDataSize    uint64
	ImageName        string `json:",omitempty"`
	ImageTimeout     time.Duration
	ImageURL         string `json:",omitempty"`
	IpAddress        net.IP
	MinimumFreeBytes uint64
	RoundupPower     uint64
} // RAW image data (length=ImageDataSize) is streamed afterwards.

type ReplaceVmImageResponse struct { // Multiple responses are sent.
	DhcpTimedOut    bool
	Final           bool // If true, this is the final response.
	ProgressMessage string
	Error           string
}

type ReplaceVmUserDataRequest struct {
	IpAddress net.IP
	Size      uint64
} // User data (length=Size) are streamed afterwards.

type ReplaceVmUserDataResponse struct {
	Error string
}

type RestoreVmFromSnapshotRequest struct {
	IpAddress         net.IP
	ForceIfNotStopped bool
}

type RestoreVmFromSnapshotResponse struct {
	Error string
}

type RestoreVmImageRequest struct {
	IpAddress net.IP
}

type RestoreVmImageResponse struct {
	Error string
}

type RestoreVmUserDataRequest struct {
	IpAddress net.IP
}

type RestoreVmUserDataResponse struct {
	Error string
}

type SnapshotVmRequest struct {
	IpAddress         net.IP
	ForceIfNotStopped bool
	RootOnly          bool
}

type SnapshotVmResponse struct {
	Error string
}

type StartVmRequest struct {
	AccessToken []byte
	DhcpTimeout time.Duration
	IpAddress   net.IP
}

type StartVmResponse struct {
	DhcpTimedOut bool
	Error        string
}

type StopVmRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type StopVmResponse struct {
	Error string
}

type State uint

type Subnet struct {
	Id                string
	IpGateway         net.IP
	IpMask            net.IP // net.IPMask can't be JSON {en,de}coded.
	DomainName        string `json:",omitempty"`
	DomainNameServers []net.IP
	Manage            bool     `json:",omitempty"`
	VlanId            uint     `json:",omitempty"`
	AllowedGroups     []string `json:",omitempty"`
	AllowedUsers      []string `json:",omitempty"`
}

type TraceVmMetadataRequest struct {
	IpAddress net.IP
}

type TraceVmMetadataResponse struct {
	Error string
} // A stream of strings (trace paths) follow.

type UpdateSubnetsRequest struct {
	Add    []Subnet
	Change []Subnet
	Delete []string
}

type UpdateSubnetsResponse struct {
	Error string
}

type VmInfo struct {
	Address            Address
	Hostname           string `json:",omitempty"`
	ImageName          string `json:",omitempty"`
	ImageURL           string `json:",omitempty"`
	MemoryInMiB        uint64
	MilliCPUs          uint
	OwnerGroups        []string `json:",omitempty"`
	OwnerUsers         []string `json:",omitempty"`
	SpreadVolumes      bool     `json:",omitempty"`
	State              State
	Tags               tags.Tags `json:",omitempty"`
	SecondaryAddresses []Address `json:",omitempty"`
	SecondarySubnetIDs []string  `json:",omitempty"`
	SubnetId           string    `json:",omitempty"`
	Uncommitted        bool      `json:",omitempty"`
	Volumes            []Volume  `json:",omitempty"`
}

type Volume struct {
	Size   uint64
	Format VolumeFormat
}

type VolumeFormat uint
<commit_msg>Add proto/hypervisor.ChangeAddressPoolRequest.AddressesToRemove field.<commit_after>package hypervisor

import (
	"net"
	"time"

	"github.com/Symantec/Dominator/lib/tags"
)

const (
	StateStarting      = 0
	StateRunning       = 1
	StateFailedToStart = 2
	StateStopping      = 3
	StateStopped       = 4
	StateDestroying    = 5
	StateMigrating     = 6

	VolumeFormatRaw   = 0
	VolumeFormatQCOW2 = 1
)

type AcknowledgeVmRequest struct {
	IpAddress net.IP
}

type AcknowledgeVmResponse struct {
	Error string
}

type Address struct {
	IpAddress  net.IP `json:",omitempty"`
	MacAddress string
}

type BecomePrimaryVmOwnerRequest struct {
	IpAddress net.IP
}

type BecomePrimaryVmOwnerResponse struct {
	Error string
}

type ChangeAddressPoolRequest struct {
	AddressesToAdd       []Address       // Will be added to free pool.
	AddressesToRemove    []Address       // Will be removed from free pool.
	MaximumFreeAddresses map[string]uint // Key: subnet ID.
}

type ChangeAddressPoolResponse struct {
	Error string
}

type ChangeOwnersRequest struct {
	OwnerGroups []string `json:",omitempty"`
	OwnerUsers  []string `json:",omitempty"`
}

type ChangeOwnersResponse struct {
	Error string
}

type ChangeVmOwnerUsersRequest struct {
	IpAddress  net.IP
	OwnerUsers []string
}

type ChangeVmOwnerUsersResponse struct {
	Error string
}

type ChangeVmTagsRequest struct {
	IpAddress net.IP
	Tags      tags.Tags
}

type ChangeVmTagsResponse struct {
	Error string
}

type CommitImportedVmRequest struct {
	IpAddress net.IP
}

type CommitImportedVmResponse struct {
	Error string
}

type CreateVmRequest struct {
	DhcpTimeout      time.Duration
	ImageDataSize    uint64
	ImageTimeout     time.Duration
	MinimumFreeBytes uint64
	RoundupPower     uint64
	SecondaryVolumes []Volume
	UserDataSize     uint64
	VmInfo
} // RAW image data (length=ImageDataSize) and user data (length=UserDataSize)
// are streamed afterwards.

type CreateVmResponse struct { // Multiple responses are sent.
	DhcpTimedOut    bool
	Final           bool // If true, this is the final response.
	IpAddress       net.IP
	ProgressMessage string
	Error           string
}

type DeleteVmVolumeRequest struct {
	AccessToken []byte
	IpAddress   net.IP
	VolumeIndex uint
}

type DeleteVmVolumeResponse struct {
	Error string
}

type DestroyVmRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type DestroyVmResponse struct {
	Error string
}

type DiscardVmAccessTokenRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type DiscardVmAccessTokenResponse struct {
	Error string
}

type DiscardVmOldImageRequest struct {
	IpAddress net.IP
}

type DiscardVmOldImageResponse struct {
	Error string
}

type DiscardVmOldUserDataRequest struct {
	IpAddress net.IP
}

type DiscardVmOldUserDataResponse struct {
	Error string
}

type DiscardVmSnapshotRequest struct {
	IpAddress net.IP
}

type DiscardVmSnapshotResponse struct {
	Error string
}

// The GetUpdates() RPC is fully streamed.
// The client may or may not send GetUpdateRequest messages to the server.
// The server sends a stream of Update messages.

type GetUpdateRequest struct{}

type Update struct {
	HaveAddressPool  bool               `json:",omitempty"`
	AddressPool      []Address          `json:",omitempty"` // Used & free.
	NumFreeAddresses map[string]uint    `json:",omitempty"` // Key: subnet ID.
	HealthStatus     string             `json:",omitempty"`
	HaveSerialNumber bool               `json:",omitempty"`
	SerialNumber     string             `json:",omitempty"`
	HaveSubnets      bool               `json:",omitempty"`
	Subnets          []Subnet           `json:",omitempty"`
	HaveVMs          bool               `json:",omitempty"`
	VMs              map[string]*VmInfo `json:",omitempty"` // Key: IP address.
}

type GetVmAccessTokenRequest struct {
	IpAddress net.IP
	Lifetime  time.Duration
}

type GetVmAccessTokenResponse struct {
	Token []byte `json:",omitempty"`
	Error string
}

type GetVmInfoRequest struct {
	IpAddress net.IP
}

type GetVmInfoResponse struct {
	VmInfo VmInfo
	Error  string
}

type GetVmUserDataRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type GetVmUserDataResponse struct {
	Error  string
	Length uint64
} // Data (length=Length) are streamed afterwards.

// The GetVmVolume() RPC is followed by the proto/rsync.GetBlocks message.

type GetVmVolumeRequest struct {
	AccessToken []byte
	IpAddress   net.IP
	VolumeIndex uint
}

type GetVmVolumeResponse struct {
	Error string
}

type ImportLocalVmRequest struct {
	VerificationCookie []byte `json:",omitempty"`
	VmInfo
	VolumeFilenames []string
}

type ImportLocalVmResponse struct {
	Error string
}

type ListVMsRequest struct {
	Sort bool
}

type ListVMsResponse struct {
	IpAddresses []net.IP
}

type ListVolumeDirectoriesRequest struct{}

type ListVolumeDirectoriesResponse struct {
	Directories []string
	Error       string
}

type MigrateVmRequest struct {
	AccessToken      []byte
	DhcpTimeout      time.Duration
	IpAddress        net.IP
	SourceHypervisor string
}

type MigrateVmResponse struct { // Multiple responses are sent.
	Error           string
	Final           bool // If true, this is the final response.
	ProgressMessage string
	RequestCommit   bool
}

type MigrateVmResponseResponse struct {
	Commit bool
}

type NetbootMachineRequest struct {
	Address                      Address
	Files                        map[string][]byte
	FilesExpiration              time.Duration
	Hostname                     string
	NumAcknowledgementsToWaitFor uint
	OfferExpiration              time.Duration
	Subnet                       *Subnet
	WaitTimeout                  time.Duration
}

type NetbootMachineResponse struct {
	Error string
}

type PrepareVmForMigrationRequest struct {
	AccessToken []byte
	Enable      bool
	IpAddress   net.IP
}

type PrepareVmForMigrationResponse struct {
	Error string
}

type ProbeVmPortRequest struct {
	IpAddress  net.IP
	PortNumber uint
	Timeout    time.Duration
}

type ProbeVmPortResponse struct {
	PortIsOpen bool
	Error      string
}

type ReplaceVmImageRequest struct {
	DhcpTimeout      time.Duration
	ImageDataSize    uint64
	ImageName        string `json:",omitempty"`
	ImageTimeout     time.Duration
	ImageURL         string `json:",omitempty"`
	IpAddress        net.IP
	MinimumFreeBytes uint64
	RoundupPower     uint64
} // RAW image data (length=ImageDataSize) is streamed afterwards.

type ReplaceVmImageResponse struct { // Multiple responses are sent.
	DhcpTimedOut    bool
	Final           bool // If true, this is the final response.
	ProgressMessage string
	Error           string
}

type ReplaceVmUserDataRequest struct {
	IpAddress net.IP
	Size      uint64
} // User data (length=Size) are streamed afterwards.

type ReplaceVmUserDataResponse struct {
	Error string
}

type RestoreVmFromSnapshotRequest struct {
	IpAddress         net.IP
	ForceIfNotStopped bool
}

type RestoreVmFromSnapshotResponse struct {
	Error string
}

type RestoreVmImageRequest struct {
	IpAddress net.IP
}

type RestoreVmImageResponse struct {
	Error string
}

type RestoreVmUserDataRequest struct {
	IpAddress net.IP
}

type RestoreVmUserDataResponse struct {
	Error string
}

type SnapshotVmRequest struct {
	IpAddress         net.IP
	ForceIfNotStopped bool
	RootOnly          bool
}

type SnapshotVmResponse struct {
	Error string
}

type StartVmRequest struct {
	AccessToken []byte
	DhcpTimeout time.Duration
	IpAddress   net.IP
}

type StartVmResponse struct {
	DhcpTimedOut bool
	Error        string
}

type StopVmRequest struct {
	AccessToken []byte
	IpAddress   net.IP
}

type StopVmResponse struct {
	Error string
}

type State uint

type Subnet struct {
	Id                string
	IpGateway         net.IP
	IpMask            net.IP // net.IPMask can't be JSON {en,de}coded.
	DomainName        string `json:",omitempty"`
	DomainNameServers []net.IP
	Manage            bool     `json:",omitempty"`
	VlanId            uint     `json:",omitempty"`
	AllowedGroups     []string `json:",omitempty"`
	AllowedUsers      []string `json:",omitempty"`
}

type TraceVmMetadataRequest struct {
	IpAddress net.IP
}

type TraceVmMetadataResponse struct {
	Error string
} // A stream of strings (trace paths) follow.

type UpdateSubnetsRequest struct {
	Add    []Subnet
	Change []Subnet
	Delete []string
}

type UpdateSubnetsResponse struct {
	Error string
}

type VmInfo struct {
	Address            Address
	Hostname           string `json:",omitempty"`
	ImageName          string `json:",omitempty"`
	ImageURL           string `json:",omitempty"`
	MemoryInMiB        uint64
	MilliCPUs          uint
	OwnerGroups        []string `json:",omitempty"`
	OwnerUsers         []string `json:",omitempty"`
	SpreadVolumes      bool     `json:",omitempty"`
	State              State
	Tags               tags.Tags `json:",omitempty"`
	SecondaryAddresses []Address `json:",omitempty"`
	SecondarySubnetIDs []string  `json:",omitempty"`
	SubnetId           string    `json:",omitempty"`
	Uncommitted        bool      `json:",omitempty"`
	Volumes            []Volume  `json:",omitempty"`
}

type Volume struct {
	Size   uint64
	Format VolumeFormat
}

type VolumeFormat uint
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gcs

import (
	"context"
	"fmt"
	"io"
	"os"
	"path"
	"sync"

	"cloud.google.com/go/storage"
	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/errorutil"
)

// UploadFunc knows how to upload into an object
type UploadFunc func(writer dataWriter) error
type destToWriter func(dest string) dataWriter

// Upload uploads all of the data in the
// uploadTargets map to GCS in parallel. The map is
// keyed on GCS path under the bucket
func Upload(bucket *storage.BucketHandle, uploadTargets map[string]UploadFunc) error {
	dtw := func(dest string) dataWriter {
		return gcsObjectWriter{bucket.Object(dest).NewWriter(context.Background())}
	}
	return upload(dtw, uploadTargets)
}

// LocalExport copies all of the data in the uploadTargets map to local files in parallel. The map
// is keyed on file path under the exportDir.
func LocalExport(exportDir string, uploadTargets map[string]UploadFunc) error {
	dtw := func(dest string) dataWriter {
		return &localFileWriter{
			filePath: path.Join(exportDir, dest),
		}
	}
	return upload(dtw, uploadTargets)
}

func upload(dtw destToWriter, uploadTargets map[string]UploadFunc) error {
	errCh := make(chan error, len(uploadTargets))
	group := &sync.WaitGroup{}
	group.Add(len(uploadTargets))
	for dest, upload := range uploadTargets {
		log := logrus.WithField("dest", dest)
		log.Info("Queued for upload")
		go func(f UploadFunc, writer dataWriter, log *logrus.Entry) {
			defer group.Done()
			if err := f(writer); err != nil {
				errCh <- err
			} else {
				log.Info("Finished upload")
			}
		}(upload, dtw(dest), log)
	}
	group.Wait()
	close(errCh)
	if len(errCh) != 0 {
		var uploadErrors []error
		for err := range errCh {
			uploadErrors = append(uploadErrors, err)
		}
		return fmt.Errorf("encountered errors during upload: %v", uploadErrors)
	}
	return nil
}

// FileUpload returns an UploadFunc which copies all
// data from the file on disk to the GCS object
func FileUpload(file string) UploadFunc {
	return FileUploadWithAttributes(file, nil)
}

// FileUploadWithMetadata returns an UploadFunc which copies all
// data from the file on disk into GCS object and also sets the provided
// metadata fields on the object.
func FileUploadWithMetadata(file string, metadata map[string]string) UploadFunc {
	return FileUploadWithAttributes(file, &storage.ObjectAttrs{Metadata: metadata})
}

// FileUploadWithAttributes returns an UploadFunc which copies all data
// from the file on disk into GCS object and also sets the provided
// attributes on the object.
func FileUploadWithAttributes(file string, attrs *storage.ObjectAttrs) UploadFunc {
	return func(writer dataWriter) error {
		reader, err := os.Open(file)
		if err != nil {
			return err
		}

		uploadErr := DataUploadWithAttributes(reader, attrs)(writer)
		if uploadErr != nil {
			uploadErr = fmt.Errorf("upload error: %v", uploadErr)
		}
		closeErr := reader.Close()
		if closeErr != nil {
			closeErr = fmt.Errorf("reader close error: %v", closeErr)
		}

		return errorutil.NewAggregate(uploadErr, closeErr)
	}
}

// DataUpload returns an UploadFunc which copies all
// data from src reader into GCS.
func DataUpload(src io.Reader) UploadFunc {
	return DataUploadWithAttributes(src, nil)
}

// DataUploadWithMetadata returns an UploadFunc which copies all
// data from src reader into GCS and also sets the provided metadata
// fields onto the object.
func DataUploadWithMetadata(src io.Reader, metadata map[string]string) UploadFunc {
	return DataUploadWithAttributes(src, &storage.ObjectAttrs{Metadata: metadata})
}

// DataUploadWithAttributes returns an UploadFunc which copies all data
// from src reader into GCS and also sets the provided attributes on
// the object.
func DataUploadWithAttributes(src io.Reader, attrs *storage.ObjectAttrs) UploadFunc {
	return func(writer dataWriter) error {
		writer.ApplyAttributes(attrs)
		_, copyErr := io.Copy(writer, src)
		if copyErr != nil {
			copyErr = fmt.Errorf("copy error: %v", copyErr)
		}
		closeErr := writer.Close()
		if closeErr != nil {
			closeErr = fmt.Errorf("writer close error: %v", closeErr)
		}
		return errorutil.NewAggregate(copyErr, closeErr)
	}
}

type dataWriter interface {
	io.WriteCloser
	ApplyAttributes(*storage.ObjectAttrs)
}

type gcsObjectWriter struct {
	*storage.Writer
}

func (w gcsObjectWriter) ApplyAttributes(attrs *storage.ObjectAttrs) {
	if attrs == nil {
		return
	}
	attrs.Name = w.Writer.ObjectAttrs.Name
	w.Writer.ObjectAttrs = *attrs
}

type localFileWriter struct {
	filePath string
	file     *os.File
}

func (w *localFileWriter) Write(b []byte) (int, error) {
	if w.file == nil {
		var err error
		w.file, err = os.OpenFile(w.filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
		if err != nil {
			return 0, fmt.Errorf("error opening %q for writing: %v", w.filePath, err)
		}
	}
	n, err := w.file.Write(b)
	if err != nil {
		return n, fmt.Errorf("error writing to %q: %v", w.filePath, err)
	}
	return n, nil
}

func (w *localFileWriter) Close() error {
	return w.file.Close()
}

// Ignore attributes when copying files locally.
func (w *localFileWriter) ApplyAttributes(_ *storage.ObjectAttrs) {}
<commit_msg>prow/pod-utils/gcs/upload: Drop unused FileUploadWithMetadata<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gcs

import (
	"context"
	"fmt"
	"io"
	"os"
	"path"
	"sync"

	"cloud.google.com/go/storage"
	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/errorutil"
)

// UploadFunc knows how to upload into an object
type UploadFunc func(writer dataWriter) error
type destToWriter func(dest string) dataWriter

// Upload uploads all of the data in the
// uploadTargets map to GCS in parallel. The map is
// keyed on GCS path under the bucket
func Upload(bucket *storage.BucketHandle, uploadTargets map[string]UploadFunc) error {
	dtw := func(dest string) dataWriter {
		return gcsObjectWriter{bucket.Object(dest).NewWriter(context.Background())}
	}
	return upload(dtw, uploadTargets)
}

// LocalExport copies all of the data in the uploadTargets map to local files in parallel. The map
// is keyed on file path under the exportDir.
func LocalExport(exportDir string, uploadTargets map[string]UploadFunc) error {
	dtw := func(dest string) dataWriter {
		return &localFileWriter{
			filePath: path.Join(exportDir, dest),
		}
	}
	return upload(dtw, uploadTargets)
}

func upload(dtw destToWriter, uploadTargets map[string]UploadFunc) error {
	errCh := make(chan error, len(uploadTargets))
	group := &sync.WaitGroup{}
	group.Add(len(uploadTargets))
	for dest, upload := range uploadTargets {
		log := logrus.WithField("dest", dest)
		log.Info("Queued for upload")
		go func(f UploadFunc, writer dataWriter, log *logrus.Entry) {
			defer group.Done()
			if err := f(writer); err != nil {
				errCh <- err
			} else {
				log.Info("Finished upload")
			}
		}(upload, dtw(dest), log)
	}
	group.Wait()
	close(errCh)
	if len(errCh) != 0 {
		var uploadErrors []error
		for err := range errCh {
			uploadErrors = append(uploadErrors, err)
		}
		return fmt.Errorf("encountered errors during upload: %v", uploadErrors)
	}
	return nil
}

// FileUpload returns an UploadFunc which copies all
// data from the file on disk to the GCS object
func FileUpload(file string) UploadFunc {
	return FileUploadWithAttributes(file, nil)
}

// FileUploadWithAttributes returns an UploadFunc which copies all data
// from the file on disk into GCS object and also sets the provided
// attributes on the object.
func FileUploadWithAttributes(file string, attrs *storage.ObjectAttrs) UploadFunc {
	return func(writer dataWriter) error {
		reader, err := os.Open(file)
		if err != nil {
			return err
		}

		uploadErr := DataUploadWithAttributes(reader, attrs)(writer)
		if uploadErr != nil {
			uploadErr = fmt.Errorf("upload error: %v", uploadErr)
		}
		closeErr := reader.Close()
		if closeErr != nil {
			closeErr = fmt.Errorf("reader close error: %v", closeErr)
		}

		return errorutil.NewAggregate(uploadErr, closeErr)
	}
}

// DataUpload returns an UploadFunc which copies all
// data from src reader into GCS.
func DataUpload(src io.Reader) UploadFunc {
	return DataUploadWithAttributes(src, nil)
}

// DataUploadWithMetadata returns an UploadFunc which copies all
// data from src reader into GCS and also sets the provided metadata
// fields onto the object.
func DataUploadWithMetadata(src io.Reader, metadata map[string]string) UploadFunc {
	return DataUploadWithAttributes(src, &storage.ObjectAttrs{Metadata: metadata})
}

// DataUploadWithAttributes returns an UploadFunc which copies all data
// from src reader into GCS and also sets the provided attributes on
// the object.
func DataUploadWithAttributes(src io.Reader, attrs *storage.ObjectAttrs) UploadFunc {
	return func(writer dataWriter) error {
		writer.ApplyAttributes(attrs)
		_, copyErr := io.Copy(writer, src)
		if copyErr != nil {
			copyErr = fmt.Errorf("copy error: %v", copyErr)
		}
		closeErr := writer.Close()
		if closeErr != nil {
			closeErr = fmt.Errorf("writer close error: %v", closeErr)
		}
		return errorutil.NewAggregate(copyErr, closeErr)
	}
}

type dataWriter interface {
	io.WriteCloser
	ApplyAttributes(*storage.ObjectAttrs)
}

type gcsObjectWriter struct {
	*storage.Writer
}

func (w gcsObjectWriter) ApplyAttributes(attrs *storage.ObjectAttrs) {
	if attrs == nil {
		return
	}
	attrs.Name = w.Writer.ObjectAttrs.Name
	w.Writer.ObjectAttrs = *attrs
}

type localFileWriter struct {
	filePath string
	file     *os.File
}

func (w *localFileWriter) Write(b []byte) (int, error) {
	if w.file == nil {
		var err error
		w.file, err = os.OpenFile(w.filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
		if err != nil {
			return 0, fmt.Errorf("error opening %q for writing: %v", w.filePath, err)
		}
	}
	n, err := w.file.Write(b)
	if err != nil {
		return n, fmt.Errorf("error writing to %q: %v", w.filePath, err)
	}
	return n, nil
}

func (w *localFileWriter) Close() error {
	return w.file.Close()
}

// Ignore attributes when copying files locally.
func (w *localFileWriter) ApplyAttributes(_ *storage.ObjectAttrs) {}
<|endoftext|>