<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go.tools/go/vcs"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
)

// Godeps describes what a package needs to be rebuilt reproducibly.
// It's the same information stored in file Godeps.
type Godeps struct {
	ImportPath string
	GoVersion  string
	Deps       []Dependency

	outerRoot string
}

// A Dependency is a specific revision of a package.
type Dependency struct {
	ImportPath string
	Comment    string `json:",omitempty"` // Description of commit, if present.
	Rev        string // VCS-specific commit ID.

	outerRoot string // dir, if present, in outer GOPATH
	repoRoot  *vcs.RepoRoot
	vcs       *VCS
}

// pkgs is the list of packages to read dependencies
func (g *Godeps) Load(pkgs []*Package) error {
	var err1 error
	var path, seen []string
	for _, p := range pkgs {
		if p.Standard {
			log.Println("ignoring stdlib package:", p.ImportPath)
			continue
		}
		if p.Error.Err != "" {
			log.Println(p.Error.Err)
			err1 = errors.New("error loading packages")
			continue
		}
		_, rr, err := VCSForImportPath(p.ImportPath)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading packages")
			continue
		}
		seen = append(seen, rr.Root)
		path = append(path, p.Deps...)
	}
	if err1 != nil {
		return err1
	}
	sort.Strings(path)
	if len(path) == 0 {
		return nil // empty list means [.] in LoadPackages; we really want []
	}
	for _, pkg := range MustLoadPackages(path...) {
		if pkg.Error.Err != "" {
			log.Println(pkg.Error.Err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if pkg.Standard {
			continue
		}
		vcs, rr, err := VCSForImportPath(pkg.ImportPath)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if contains(seen, rr.Root) {
			continue
		}
		seen = append(seen, rr.Root)
		var id string
		id, err = vcs.identify(pkg.Dir)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if vcs.isDirty(pkg.Dir) {
			log.Println("dirty working tree:", pkg.Dir)
			err1 = errors.New("error loading dependencies")
			continue
		}
		comment := vcs.describe(pkg.Dir, id)
		g.Deps = append(g.Deps, Dependency{
			ImportPath: pkg.ImportPath,
			Rev:        id,
			Comment:    comment,
		})
	}
	return err1
}

func ReadGodeps(path string) (*Godeps, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	g := new(Godeps)
	err = json.NewDecoder(f).Decode(g)
	if err != nil {
		return nil, err
	}
	err = g.loadGoList()
	if err != nil {
		return nil, err
	}

	for i := range g.Deps {
		d := &g.Deps[i]
		d.vcs, d.repoRoot, err = VCSForImportPath(d.ImportPath)
		if err != nil {
			return nil, err
		}
	}
	return g, nil
}

func (g *Godeps) loadGoList() error {
	a := []string{g.ImportPath}
	for _, d := range g.Deps {
		a = append(a, d.ImportPath)
	}
	ps, err := LoadPackages(a...)
	if err != nil {
		return err
	}
	g.outerRoot = ps[0].Root
	for i, p := range ps[1:] {
		g.Deps[i].outerRoot = p.Root
	}
	return nil
}

func (g *Godeps) WriteTo(w io.Writer) (int, error) {
	b, err := json.MarshalIndent(g, "", "\t")
	if err != nil {
		return 0, err
	}
	return w.Write(append(b, '\n'))
}

// Returns a path to the local copy of d's repository.
// E.g.
//
//   ImportPath             RepoPath
//   github.com/kr/s3       $spool/github.com/kr/s3
//   github.com/lib/pq/oid  $spool/github.com/lib/pq
func (d Dependency) RepoPath() string {
	return filepath.Join(spool, "repo", d.repoRoot.Root)
}

// Returns a URL for the remote copy of the repository.
func (d Dependency) RemoteURL() string {
	return d.repoRoot.Repo
}

// Returns the url of a local disk clone of the repo, if any.
func (d Dependency) FastRemotePath() string {
	if d.outerRoot != "" {
		return d.outerRoot + "/src/" + d.repoRoot.Root
	}
	return ""
}

// Returns a path to the checked-out copy of d's commit.
func (d Dependency) Workdir() string {
	return filepath.Join(d.Gopath(), "src", d.ImportPath)
}

// Returns a path to the checked-out copy of d's repo root.
func (d Dependency) WorkdirRoot() string {
	return filepath.Join(d.Gopath(), "src", d.repoRoot.Root)
}

// Returns a path to a parent of Workdir such that using
// Gopath in GOPATH makes d available to the go tool.
func (d Dependency) Gopath() string {
	return filepath.Join(spool, "rev", d.Rev[:2], d.Rev[2:])
}

// Creates an empty repo in d.RepoPath().
func (d Dependency) CreateRepo(fastRemote, mainRemote string) error {
	if err := os.MkdirAll(d.RepoPath(), 0777); err != nil {
		return err
	}
	if err := d.vcs.create(d.RepoPath()); err != nil {
		return err
	}
	if err := d.link(fastRemote, d.FastRemotePath()); err != nil {
		return err
	}
	return d.link(mainRemote, d.RemoteURL())
}

func (d Dependency) link(remote, url string) error {
	return d.vcs.link(d.RepoPath(), remote, url)
}

func (d Dependency) fetchAndCheckout(remote string) error {
	if err := d.fetch(remote); err != nil {
		return fmt.Errorf("fetch: %s", err)
	}
	if err := d.checkout(); err != nil {
		return fmt.Errorf("checkout: %s", err)
	}
	return nil
}

func (d Dependency) fetch(remote string) error {
	return d.vcs.fetch(d.RepoPath(), remote)
}

func (d Dependency) checkout() error {
	dir := d.WorkdirRoot()
	if exists(dir) {
		return nil
	}
	if !d.vcs.exists(d.RepoPath(), d.Rev) {
		return fmt.Errorf("unknown rev %s", d.Rev)
	}
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}
	return d.vcs.checkout(dir, d.Rev, d.RepoPath())
}

func contains(a []string, s string) bool {
	for _, p := range a {
		if s == p {
			return true
		}
	}
	return false
}

func mustGoVersion() string {
	cmd := exec.Command("go", "version")
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	return string(bytes.TrimSpace(out))
}
<commit_msg>refactor<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go.tools/go/vcs"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
)

// Godeps describes what a package needs to be rebuilt reproducibly.
// It's the same information stored in file Godeps.
type Godeps struct {
	ImportPath string
	GoVersion  string
	Deps       []Dependency

	outerRoot string
}

// A Dependency is a specific revision of a package.
type Dependency struct {
	ImportPath string
	Comment    string `json:",omitempty"` // Description of commit, if present.
	Rev        string // VCS-specific commit ID.

	outerRoot string // dir, if present, in outer GOPATH
	repoRoot  *vcs.RepoRoot
	vcs       *VCS
}

// pkgs is the list of packages to read dependencies
func (g *Godeps) Load(pkgs []*Package) error {
	var err1 error
	var path, seen []string
	for _, p := range pkgs {
		if p.Standard {
			log.Println("ignoring stdlib package:", p.ImportPath)
			continue
		}
		if p.Error.Err != "" {
			log.Println(p.Error.Err)
			err1 = errors.New("error loading packages")
			continue
		}
		_, rr, err := VCSForImportPath(p.ImportPath)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading packages")
			continue
		}
		seen = append(seen, rr.Root)
		path = append(path, p.Deps...)
	}
	if err1 != nil {
		return err1
	}
	sort.Strings(path)
	if len(path) == 0 {
		return nil // empty list means [.] in LoadPackages; we really want []
	}
	for _, pkg := range MustLoadPackages(path...) {
		if pkg.Error.Err != "" {
			log.Println(pkg.Error.Err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if pkg.Standard {
			continue
		}
		vcs, rr, err := VCSForImportPath(pkg.ImportPath)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if contains(seen, rr.Root) {
			continue
		}
		seen = append(seen, rr.Root)
		id, err := vcs.identify(pkg.Dir)
		if err != nil {
			log.Println(err)
			err1 = errors.New("error loading dependencies")
			continue
		}
		if vcs.isDirty(pkg.Dir) {
			log.Println("dirty working tree:", pkg.Dir)
			err1 = errors.New("error loading dependencies")
			continue
		}
		comment := vcs.describe(pkg.Dir, id)
		g.Deps = append(g.Deps, Dependency{
			ImportPath: pkg.ImportPath,
			Rev:        id,
			Comment:    comment,
		})
	}
	return err1
}

func ReadGodeps(path string) (*Godeps, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	g := new(Godeps)
	err = json.NewDecoder(f).Decode(g)
	if err != nil {
		return nil, err
	}
	err = g.loadGoList()
	if err != nil {
		return nil, err
	}

	for i := range g.Deps {
		d := &g.Deps[i]
		d.vcs, d.repoRoot, err = VCSForImportPath(d.ImportPath)
		if err != nil {
			return nil, err
		}
	}
	return g, nil
}

func (g *Godeps) loadGoList() error {
	a := []string{g.ImportPath}
	for _, d := range g.Deps {
		a = append(a, d.ImportPath)
	}
	ps, err := LoadPackages(a...)
	if err != nil {
		return err
	}
	g.outerRoot = ps[0].Root
	for i, p := range ps[1:] {
		g.Deps[i].outerRoot = p.Root
	}
	return nil
}

func (g *Godeps) WriteTo(w io.Writer) (int, error) {
	b, err := json.MarshalIndent(g, "", "\t")
	if err != nil {
		return 0, err
	}
	return w.Write(append(b, '\n'))
}

// Returns a path to the local copy of d's repository.
// E.g.
//
//   ImportPath             RepoPath
//   github.com/kr/s3       $spool/github.com/kr/s3
//   github.com/lib/pq/oid  $spool/github.com/lib/pq
func (d Dependency) RepoPath() string {
	return filepath.Join(spool, "repo", d.repoRoot.Root)
}

// Returns a URL for the remote copy of the repository.
func (d Dependency) RemoteURL() string {
	return d.repoRoot.Repo
}

// Returns the url of a local disk clone of the repo, if any.
func (d Dependency) FastRemotePath() string {
	if d.outerRoot != "" {
		return d.outerRoot + "/src/" + d.repoRoot.Root
	}
	return ""
}

// Returns a path to the checked-out copy of d's commit.
func (d Dependency) Workdir() string {
	return filepath.Join(d.Gopath(), "src", d.ImportPath)
}

// Returns a path to the checked-out copy of d's repo root.
func (d Dependency) WorkdirRoot() string {
	return filepath.Join(d.Gopath(), "src", d.repoRoot.Root)
}

// Returns a path to a parent of Workdir such that using
// Gopath in GOPATH makes d available to the go tool.
func (d Dependency) Gopath() string {
	return filepath.Join(spool, "rev", d.Rev[:2], d.Rev[2:])
}

// Creates an empty repo in d.RepoPath().
func (d Dependency) CreateRepo(fastRemote, mainRemote string) error {
	if err := os.MkdirAll(d.RepoPath(), 0777); err != nil {
		return err
	}
	if err := d.vcs.create(d.RepoPath()); err != nil {
		return err
	}
	if err := d.link(fastRemote, d.FastRemotePath()); err != nil {
		return err
	}
	return d.link(mainRemote, d.RemoteURL())
}

func (d Dependency) link(remote, url string) error {
	return d.vcs.link(d.RepoPath(), remote, url)
}

func (d Dependency) fetchAndCheckout(remote string) error {
	if err := d.fetch(remote); err != nil {
		return fmt.Errorf("fetch: %s", err)
	}
	if err := d.checkout(); err != nil {
		return fmt.Errorf("checkout: %s", err)
	}
	return nil
}

func (d Dependency) fetch(remote string) error {
	return d.vcs.fetch(d.RepoPath(), remote)
}

func (d Dependency) checkout() error {
	dir := d.WorkdirRoot()
	if exists(dir) {
		return nil
	}
	if !d.vcs.exists(d.RepoPath(), d.Rev) {
		return fmt.Errorf("unknown rev %s", d.Rev)
	}
	if err := os.MkdirAll(dir, 0777); err != nil {
		return err
	}
	return d.vcs.checkout(dir, d.Rev, d.RepoPath())
}

func contains(a []string, s string) bool {
	for _, p := range a {
		if s == p {
			return true
		}
	}
	return false
}

func mustGoVersion() string {
	cmd := exec.Command("go", "version")
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	return string(bytes.TrimSpace(out))
}
<|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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.

// Matchers use a tri-state logic in order to make the semantics of matchers
// that wrap other matchers make more sense. The constants below represent the
// three values that a matcher may return.
const (
	// MATCH_FALSE indicates that the supplied value didn't match. For example,
	// IsNil would return this when presented with any non-nil value, and
	// GreaterThan(17) would return this when presented with 16.
	MATCH_FALSE = 0

	// MATCH_TRUE indicates that the supplied value did match. For example, IsNil
	// would return this when presented with nil, and GreaterThan(17) would
	// return this when presented with 19.
	MATCH_TRUE = 1

	// MATCH_UNDEFINED indicates that the matcher doesn't process values of the
	// supplied type, or otherwise doesn't know how to handle the value. This is
	// akin to returning MATCH_FALSE, except that wrapper matchers should
	// propagagate undefined values.
	//
	// For example, if GreaterThan(17) returned MATCH_FALSE for the value "taco",
	// then Not(GreaterThan(17)) would return MATCH_TRUE. This is technically
	// correct, but is surprising and may mask failures where the wrong sort of
	// matcher is accidentally used. Instead, GreaterThan(17) can return
	// MATCH_UNDEFINED, which will be propagated by Not().
	MATCH_UNDEFINED = -1
)

// A MatchResult is an integer equal to one of the MATCH_* constants above.
type MatchResult int

// A Matcher is some predicate implicitly defining a set of values that it
// matches. For example, GreaterThan(17) matches all numeric values greater
// than 17, and HasSubstr("taco") matches all strings with the substring
// "taco".
type Matcher interface {
	// Matches returns a MatchResult indicating whether the supplied value
	// belongs to the set defined by the matcher.
	//
	// If the result is MATCH_FALSE or MATCH_UNDEFINED, it additionally returns
	// an error string describing why the value doesn't match. Error strings are
	// relative clauses that are suitable for being placed after the value. For
	// example, a predicate that matches strings with a particular substring may,
	// when presented with a numerical value, return the following string:
	//
	//     "which is not a string"
	//
	// Then the failure message may look like:
	//
	//     Expected: is a string with substring "taco"
	//     Actual:   17, which is not a string
	//
	func Matches(val interface{}) (result MatchResult, error string)

	// Description returns a string describing the property that values matching
	// this matcher have, as a verb phrase where the subject is the value. For
	// example, "is greather than 17" or "is a string with substring "taco"".
	func Description() string
}
<commit_msg>Fixed some errors.<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 ogletest

// Matchers use a tri-state logic in order to make the semantics of matchers
// that wrap other matchers make more sense. The constants below represent the
// three values that a matcher may return.
const (
	// MATCH_FALSE indicates that the supplied value didn't match. For example,
	// IsNil would return this when presented with any non-nil value, and
	// GreaterThan(17) would return this when presented with 16.
	MATCH_FALSE = 0

	// MATCH_TRUE indicates that the supplied value did match. For example, IsNil
	// would return this when presented with nil, and GreaterThan(17) would
	// return this when presented with 19.
	MATCH_TRUE = 1

	// MATCH_UNDEFINED indicates that the matcher doesn't process values of the
	// supplied type, or otherwise doesn't know how to handle the value. This is
	// akin to returning MATCH_FALSE, except that wrapper matchers should
	// propagagate undefined values.
	//
	// For example, if GreaterThan(17) returned MATCH_FALSE for the value "taco",
	// then Not(GreaterThan(17)) would return MATCH_TRUE. This is technically
	// correct, but is surprising and may mask failures where the wrong sort of
	// matcher is accidentally used. Instead, GreaterThan(17) can return
	// MATCH_UNDEFINED, which will be propagated by Not().
	MATCH_UNDEFINED = -1
)

// A MatchResult is an integer equal to one of the MATCH_* constants above.
type MatchResult int

// A Matcher is some predicate implicitly defining a set of values that it
// matches. For example, GreaterThan(17) matches all numeric values greater
// than 17, and HasSubstr("taco") matches all strings with the substring
// "taco".
type Matcher interface {
	// Matches returns a MatchResult indicating whether the supplied value
	// belongs to the set defined by the matcher.
	//
	// If the result is MATCH_FALSE or MATCH_UNDEFINED, it additionally returns
	// an error string describing why the value doesn't match. Error strings are
	// relative clauses that are suitable for being placed after the value. For
	// example, a predicate that matches strings with a particular substring may,
	// when presented with a numerical value, return the following string:
	//
	//     "which is not a string"
	//
	// Then the failure message may look like:
	//
	//     Expected: is a string with substring "taco"
	//     Actual:   17, which is not a string
	//
	Matches(val interface{}) (result MatchResult, error string)

	// Description returns a string describing the property that values matching
	// this matcher have, as a verb phrase where the subject is the value. For
	// example, "is greather than 17" or "is a string with substring "taco"".
	Description() string
}
<|endoftext|>


<commit_before>// Package dht implements a distributed hash table that satisfies the ipfs routing
// interface. This DHT is modeled after kademlia with S/Kademlia modifications.
package dht

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	opts "github.com/libp2p/go-libp2p-kad-dht/opts"
	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	providers "github.com/libp2p/go-libp2p-kad-dht/providers"

	proto "github.com/gogo/protobuf/proto"
	cid "github.com/ipfs/go-cid"
	ds "github.com/ipfs/go-datastore"
	logging "github.com/ipfs/go-log"
	goprocess "github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	ci "github.com/libp2p/go-libp2p-crypto"
	host "github.com/libp2p/go-libp2p-host"
	kb "github.com/libp2p/go-libp2p-kbucket"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	protocol "github.com/libp2p/go-libp2p-protocol"
	record "github.com/libp2p/go-libp2p-record"
	recpb "github.com/libp2p/go-libp2p-record/pb"
	routing "github.com/libp2p/go-libp2p-routing"
	base32 "github.com/whyrusleeping/base32"
)

var log = logging.Logger("dht")

var ProtocolDHT protocol.ID = "/ipfs/kad/1.0.0"
var ProtocolDHTOld protocol.ID = "/ipfs/dht"

// NumBootstrapQueries defines the number of random dht queries to do to
// collect members of the routing table.
const NumBootstrapQueries = 5

// IpfsDHT is an implementation of Kademlia with S/Kademlia modifications.
// It is used to implement the base IpfsRouting module.
type IpfsDHT struct {
	host      host.Host        // the network services we need
	self      peer.ID          // Local peer (yourself)
	peerstore pstore.Peerstore // Peer Registry

	datastore ds.Datastore // Local data

	routingTable *kb.RoutingTable // Array of routing tables for differently distanced nodes
	providers    *providers.ProviderManager

	birth time.Time // When this peer started up

	Validator record.Validator

	ctx  context.Context
	proc goprocess.Process

	strmap map[peer.ID]*messageSender
	smlk   sync.Mutex

	plk sync.Mutex
}

// New creates a new DHT with the specified host and options.
func New(ctx context.Context, h host.Host, options ...opts.Option) (*IpfsDHT, error) {
	var cfg opts.Options
	if err := cfg.Apply(append([]opts.Option{opts.Defaults}, options...)...); err != nil {
		return nil, err
	}
	dht := makeDHT(ctx, h, cfg.Datastore)

	// register for network notifs.
	dht.host.Network().Notify((*netNotifiee)(dht))

	dht.proc = goprocessctx.WithContextAndTeardown(ctx, func() error {
		// remove ourselves from network notifs.
		dht.host.Network().StopNotify((*netNotifiee)(dht))
		return nil
	})

	dht.proc.AddChild(dht.providers.Process())
	dht.Validator = cfg.Validator

	if !cfg.Client {
		h.SetStreamHandler(ProtocolDHT, dht.handleNewStream)
		h.SetStreamHandler(ProtocolDHTOld, dht.handleNewStream)
	}
	return dht, nil
}

// NewDHT creates a new DHT object with the given peer as the 'local' host.
// IpfsDHT's initialized with this function will respond to DHT requests,
// whereas IpfsDHT's initialized with NewDHTClient will not.
func NewDHT(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	dht, err := New(ctx, h, opts.Datastore(dstore))
	if err != nil {
		panic(err)
	}
	return dht
}

// NewDHTClient creates a new DHT object with the given peer as the 'local'
// host. IpfsDHT clients initialized with this function will not respond to DHT
// requests. If you need a peer to respond to DHT requests, use NewDHT instead.
// NewDHTClient creates a new DHT object with the given peer as the 'local' host
func NewDHTClient(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	dht, err := New(ctx, h, opts.Datastore(dstore), opts.Client(true))
	if err != nil {
		panic(err)
	}
	return dht
}

func makeDHT(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	rt := kb.NewRoutingTable(KValue, kb.ConvertPeerID(h.ID()), time.Minute, h.Peerstore())

	cmgr := h.ConnManager()
	rt.PeerAdded = func(p peer.ID) {
		cmgr.TagPeer(p, "kbucket", 5)
	}
	rt.PeerRemoved = func(p peer.ID) {
		cmgr.UntagPeer(p, "kbucket")
	}

	return &IpfsDHT{
		datastore:    dstore,
		self:         h.ID(),
		peerstore:    h.Peerstore(),
		host:         h,
		strmap:       make(map[peer.ID]*messageSender),
		ctx:          ctx,
		providers:    providers.NewProviderManager(ctx, h.ID(), dstore),
		birth:        time.Now(),
		routingTable: rt,
	}
}

// putValueToPeer stores the given key/value pair at the peer 'p'
func (dht *IpfsDHT) putValueToPeer(ctx context.Context, p peer.ID,
	key string, rec *recpb.Record) error {

	pmes := pb.NewMessage(pb.Message_PUT_VALUE, key, 0)
	pmes.Record = rec
	rpmes, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), key)
		fallthrough
	default:
		return err
	case nil:
		break
	}

	if err != nil {
		return err
	}

	if !bytes.Equal(rpmes.GetRecord().Value, pmes.GetRecord().Value) {
		return errors.New("value not put correctly")
	}
	return nil
}

var errInvalidRecord = errors.New("received invalid record")

// getValueOrPeers queries a particular peer p for the value for
// key. It returns either the value or a list of closer peers.
// NOTE: It will update the dht's peerstore with any new addresses
// it finds for the given peer.
func (dht *IpfsDHT) getValueOrPeers(ctx context.Context, p peer.ID, key string) (*recpb.Record, []*pstore.PeerInfo, error) {

	pmes, err := dht.getValueSingle(ctx, p, key)
	if err != nil {
		return nil, nil, err
	}

	// Perhaps we were given closer peers
	peers := pb.PBPeersToPeerInfos(pmes.GetCloserPeers())

	if record := pmes.GetRecord(); record != nil {
		// Success! We were given the value
		log.Debug("getValueOrPeers: got value")

		// make sure record is valid.
		err = dht.Validator.Validate(record.GetKey(), record.GetValue())
		if err != nil {
			log.Info("Received invalid record! (discarded)")
			// return a sentinal to signify an invalid record was received
			err = errInvalidRecord
			record = new(recpb.Record)
		}
		return record, peers, err
	}

	if len(peers) > 0 {
		log.Debug("getValueOrPeers: peers")
		return nil, peers, nil
	}

	log.Warning("getValueOrPeers: routing.ErrNotFound")
	return nil, nil, routing.ErrNotFound
}

// getValueSingle simply performs the get value RPC with the given parameters
func (dht *IpfsDHT) getValueSingle(ctx context.Context, p peer.ID, key string) (*pb.Message, error) {
	meta := logging.LoggableMap{
		"key":  key,
		"peer": p,
	}

	eip := log.EventBegin(ctx, "getValueSingle", meta)
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_GET_VALUE, key, 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), key)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

// getLocal attempts to retrieve the value from the datastore
func (dht *IpfsDHT) getLocal(key string) (*recpb.Record, error) {
	log.Debugf("getLocal %s", key)

	v, err := dht.datastore.Get(mkDsKey(key))
	if err != nil {
		return nil, err
	}
	log.Debugf("found %s in local datastore")

	byt, ok := v.([]byte)
	if !ok {
		return nil, errors.New("value stored in datastore not []byte")
	}
	rec := new(recpb.Record)
	err = proto.Unmarshal(byt, rec)
	if err != nil {
		return nil, err
	}

	err = dht.Validator.Validate(rec.GetKey(), rec.GetValue())
	if err != nil {
		log.Debugf("local record verify failed: %s (discarded)", err)
		return nil, err
	}

	return rec, nil
}

// getOwnPrivateKey attempts to load the local peers private
// key from the peerstore.
func (dht *IpfsDHT) getOwnPrivateKey() (ci.PrivKey, error) {
	sk := dht.peerstore.PrivKey(dht.self)
	if sk == nil {
		log.Warningf("%s dht cannot get own private key!", dht.self)
		return nil, fmt.Errorf("cannot get private key to sign record!")
	}
	return sk, nil
}

// putLocal stores the key value pair in the datastore
func (dht *IpfsDHT) putLocal(key string, rec *recpb.Record) error {
	data, err := proto.Marshal(rec)
	if err != nil {
		return err
	}

	return dht.datastore.Put(mkDsKey(key), data)
}

// Update signals the routingTable to Update its last-seen status
// on the given peer.
func (dht *IpfsDHT) Update(ctx context.Context, p peer.ID) {
	log.Event(ctx, "updatePeer", p)
	dht.routingTable.Update(p)
}

// FindLocal looks for a peer with a given ID connected to this dht and returns the peer and the table it was found in.
func (dht *IpfsDHT) FindLocal(id peer.ID) pstore.PeerInfo {
	p := dht.routingTable.Find(id)
	if p != "" {
		return dht.peerstore.PeerInfo(p)
	}
	return pstore.PeerInfo{}
}

// findPeerSingle asks peer 'p' if they know where the peer with id 'id' is
func (dht *IpfsDHT) findPeerSingle(ctx context.Context, p peer.ID, id peer.ID) (*pb.Message, error) {
	eip := log.EventBegin(ctx, "findPeerSingle",
		logging.LoggableMap{
			"peer":   p,
			"target": id,
		})
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_FIND_NODE, string(id), 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), id)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

func (dht *IpfsDHT) findProvidersSingle(ctx context.Context, p peer.ID, key *cid.Cid) (*pb.Message, error) {
	eip := log.EventBegin(ctx, "findProvidersSingle", p, key)
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_GET_PROVIDERS, key.KeyString(), 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), key)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

// nearestPeersToQuery returns the routing tables closest peers.
func (dht *IpfsDHT) nearestPeersToQuery(pmes *pb.Message, count int) []peer.ID {
	closer := dht.routingTable.NearestPeers(kb.ConvertKey(pmes.GetKey()), count)
	return closer
}

// betterPeerToQuery returns nearestPeersToQuery, but iff closer than self.
func (dht *IpfsDHT) betterPeersToQuery(pmes *pb.Message, p peer.ID, count int) []peer.ID {
	closer := dht.nearestPeersToQuery(pmes, count)

	// no node? nil
	if closer == nil {
		log.Warning("no closer peers to send:", p)
		return nil
	}

	filtered := make([]peer.ID, 0, len(closer))
	for _, clp := range closer {

		// == to self? thats bad
		if clp == dht.self {
			log.Warning("attempted to return self! this shouldn't happen...")
			return nil
		}
		// Dont send a peer back themselves
		if clp == p {
			continue
		}

		filtered = append(filtered, clp)
	}

	// ok seems like closer nodes
	return filtered
}

// Context return dht's context
func (dht *IpfsDHT) Context() context.Context {
	return dht.ctx
}

// Process return dht's process
func (dht *IpfsDHT) Process() goprocess.Process {
	return dht.proc
}

// Close calls Process Close
func (dht *IpfsDHT) Close() error {
	return dht.proc.Close()
}

func mkDsKey(s string) ds.Key {
	return ds.NewKey(base32.RawStdEncoding.EncodeToString([]byte(s)))
}
<commit_msg>remove redundant error check<commit_after>// Package dht implements a distributed hash table that satisfies the ipfs routing
// interface. This DHT is modeled after kademlia with S/Kademlia modifications.
package dht

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	opts "github.com/libp2p/go-libp2p-kad-dht/opts"
	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
	providers "github.com/libp2p/go-libp2p-kad-dht/providers"

	proto "github.com/gogo/protobuf/proto"
	cid "github.com/ipfs/go-cid"
	ds "github.com/ipfs/go-datastore"
	logging "github.com/ipfs/go-log"
	goprocess "github.com/jbenet/goprocess"
	goprocessctx "github.com/jbenet/goprocess/context"
	ci "github.com/libp2p/go-libp2p-crypto"
	host "github.com/libp2p/go-libp2p-host"
	kb "github.com/libp2p/go-libp2p-kbucket"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	protocol "github.com/libp2p/go-libp2p-protocol"
	record "github.com/libp2p/go-libp2p-record"
	recpb "github.com/libp2p/go-libp2p-record/pb"
	routing "github.com/libp2p/go-libp2p-routing"
	base32 "github.com/whyrusleeping/base32"
)

var log = logging.Logger("dht")

var ProtocolDHT protocol.ID = "/ipfs/kad/1.0.0"
var ProtocolDHTOld protocol.ID = "/ipfs/dht"

// NumBootstrapQueries defines the number of random dht queries to do to
// collect members of the routing table.
const NumBootstrapQueries = 5

// IpfsDHT is an implementation of Kademlia with S/Kademlia modifications.
// It is used to implement the base IpfsRouting module.
type IpfsDHT struct {
	host      host.Host        // the network services we need
	self      peer.ID          // Local peer (yourself)
	peerstore pstore.Peerstore // Peer Registry

	datastore ds.Datastore // Local data

	routingTable *kb.RoutingTable // Array of routing tables for differently distanced nodes
	providers    *providers.ProviderManager

	birth time.Time // When this peer started up

	Validator record.Validator

	ctx  context.Context
	proc goprocess.Process

	strmap map[peer.ID]*messageSender
	smlk   sync.Mutex

	plk sync.Mutex
}

// New creates a new DHT with the specified host and options.
func New(ctx context.Context, h host.Host, options ...opts.Option) (*IpfsDHT, error) {
	var cfg opts.Options
	if err := cfg.Apply(append([]opts.Option{opts.Defaults}, options...)...); err != nil {
		return nil, err
	}
	dht := makeDHT(ctx, h, cfg.Datastore)

	// register for network notifs.
	dht.host.Network().Notify((*netNotifiee)(dht))

	dht.proc = goprocessctx.WithContextAndTeardown(ctx, func() error {
		// remove ourselves from network notifs.
		dht.host.Network().StopNotify((*netNotifiee)(dht))
		return nil
	})

	dht.proc.AddChild(dht.providers.Process())
	dht.Validator = cfg.Validator

	if !cfg.Client {
		h.SetStreamHandler(ProtocolDHT, dht.handleNewStream)
		h.SetStreamHandler(ProtocolDHTOld, dht.handleNewStream)
	}
	return dht, nil
}

// NewDHT creates a new DHT object with the given peer as the 'local' host.
// IpfsDHT's initialized with this function will respond to DHT requests,
// whereas IpfsDHT's initialized with NewDHTClient will not.
func NewDHT(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	dht, err := New(ctx, h, opts.Datastore(dstore))
	if err != nil {
		panic(err)
	}
	return dht
}

// NewDHTClient creates a new DHT object with the given peer as the 'local'
// host. IpfsDHT clients initialized with this function will not respond to DHT
// requests. If you need a peer to respond to DHT requests, use NewDHT instead.
// NewDHTClient creates a new DHT object with the given peer as the 'local' host
func NewDHTClient(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	dht, err := New(ctx, h, opts.Datastore(dstore), opts.Client(true))
	if err != nil {
		panic(err)
	}
	return dht
}

func makeDHT(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT {
	rt := kb.NewRoutingTable(KValue, kb.ConvertPeerID(h.ID()), time.Minute, h.Peerstore())

	cmgr := h.ConnManager()
	rt.PeerAdded = func(p peer.ID) {
		cmgr.TagPeer(p, "kbucket", 5)
	}
	rt.PeerRemoved = func(p peer.ID) {
		cmgr.UntagPeer(p, "kbucket")
	}

	return &IpfsDHT{
		datastore:    dstore,
		self:         h.ID(),
		peerstore:    h.Peerstore(),
		host:         h,
		strmap:       make(map[peer.ID]*messageSender),
		ctx:          ctx,
		providers:    providers.NewProviderManager(ctx, h.ID(), dstore),
		birth:        time.Now(),
		routingTable: rt,
	}
}

// putValueToPeer stores the given key/value pair at the peer 'p'
func (dht *IpfsDHT) putValueToPeer(ctx context.Context, p peer.ID,
	key string, rec *recpb.Record) error {

	pmes := pb.NewMessage(pb.Message_PUT_VALUE, key, 0)
	pmes.Record = rec
	rpmes, err := dht.sendRequest(ctx, p, pmes)
	if err != nil {
		if err == ErrReadTimeout {
			log.Warningf("read timeout: %s %s", p.Pretty(), key)
		}
		return err
	}

	if !bytes.Equal(rpmes.GetRecord().Value, pmes.GetRecord().Value) {
		return errors.New("value not put correctly")
	}
	return nil
}

var errInvalidRecord = errors.New("received invalid record")

// getValueOrPeers queries a particular peer p for the value for
// key. It returns either the value or a list of closer peers.
// NOTE: It will update the dht's peerstore with any new addresses
// it finds for the given peer.
func (dht *IpfsDHT) getValueOrPeers(ctx context.Context, p peer.ID, key string) (*recpb.Record, []*pstore.PeerInfo, error) {

	pmes, err := dht.getValueSingle(ctx, p, key)
	if err != nil {
		return nil, nil, err
	}

	// Perhaps we were given closer peers
	peers := pb.PBPeersToPeerInfos(pmes.GetCloserPeers())

	if record := pmes.GetRecord(); record != nil {
		// Success! We were given the value
		log.Debug("getValueOrPeers: got value")

		// make sure record is valid.
		err = dht.Validator.Validate(record.GetKey(), record.GetValue())
		if err != nil {
			log.Info("Received invalid record! (discarded)")
			// return a sentinal to signify an invalid record was received
			err = errInvalidRecord
			record = new(recpb.Record)
		}
		return record, peers, err
	}

	if len(peers) > 0 {
		log.Debug("getValueOrPeers: peers")
		return nil, peers, nil
	}

	log.Warning("getValueOrPeers: routing.ErrNotFound")
	return nil, nil, routing.ErrNotFound
}

// getValueSingle simply performs the get value RPC with the given parameters
func (dht *IpfsDHT) getValueSingle(ctx context.Context, p peer.ID, key string) (*pb.Message, error) {
	meta := logging.LoggableMap{
		"key":  key,
		"peer": p,
	}

	eip := log.EventBegin(ctx, "getValueSingle", meta)
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_GET_VALUE, key, 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), key)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

// getLocal attempts to retrieve the value from the datastore
func (dht *IpfsDHT) getLocal(key string) (*recpb.Record, error) {
	log.Debugf("getLocal %s", key)

	v, err := dht.datastore.Get(mkDsKey(key))
	if err != nil {
		return nil, err
	}
	log.Debugf("found %s in local datastore")

	byt, ok := v.([]byte)
	if !ok {
		return nil, errors.New("value stored in datastore not []byte")
	}
	rec := new(recpb.Record)
	err = proto.Unmarshal(byt, rec)
	if err != nil {
		return nil, err
	}

	err = dht.Validator.Validate(rec.GetKey(), rec.GetValue())
	if err != nil {
		log.Debugf("local record verify failed: %s (discarded)", err)
		return nil, err
	}

	return rec, nil
}

// getOwnPrivateKey attempts to load the local peers private
// key from the peerstore.
func (dht *IpfsDHT) getOwnPrivateKey() (ci.PrivKey, error) {
	sk := dht.peerstore.PrivKey(dht.self)
	if sk == nil {
		log.Warningf("%s dht cannot get own private key!", dht.self)
		return nil, fmt.Errorf("cannot get private key to sign record!")
	}
	return sk, nil
}

// putLocal stores the key value pair in the datastore
func (dht *IpfsDHT) putLocal(key string, rec *recpb.Record) error {
	data, err := proto.Marshal(rec)
	if err != nil {
		return err
	}

	return dht.datastore.Put(mkDsKey(key), data)
}

// Update signals the routingTable to Update its last-seen status
// on the given peer.
func (dht *IpfsDHT) Update(ctx context.Context, p peer.ID) {
	log.Event(ctx, "updatePeer", p)
	dht.routingTable.Update(p)
}

// FindLocal looks for a peer with a given ID connected to this dht and returns the peer and the table it was found in.
func (dht *IpfsDHT) FindLocal(id peer.ID) pstore.PeerInfo {
	p := dht.routingTable.Find(id)
	if p != "" {
		return dht.peerstore.PeerInfo(p)
	}
	return pstore.PeerInfo{}
}

// findPeerSingle asks peer 'p' if they know where the peer with id 'id' is
func (dht *IpfsDHT) findPeerSingle(ctx context.Context, p peer.ID, id peer.ID) (*pb.Message, error) {
	eip := log.EventBegin(ctx, "findPeerSingle",
		logging.LoggableMap{
			"peer":   p,
			"target": id,
		})
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_FIND_NODE, string(id), 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), id)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

func (dht *IpfsDHT) findProvidersSingle(ctx context.Context, p peer.ID, key *cid.Cid) (*pb.Message, error) {
	eip := log.EventBegin(ctx, "findProvidersSingle", p, key)
	defer eip.Done()

	pmes := pb.NewMessage(pb.Message_GET_PROVIDERS, key.KeyString(), 0)
	resp, err := dht.sendRequest(ctx, p, pmes)
	switch err {
	case nil:
		return resp, nil
	case ErrReadTimeout:
		log.Warningf("read timeout: %s %s", p.Pretty(), key)
		fallthrough
	default:
		eip.SetError(err)
		return nil, err
	}
}

// nearestPeersToQuery returns the routing tables closest peers.
func (dht *IpfsDHT) nearestPeersToQuery(pmes *pb.Message, count int) []peer.ID {
	closer := dht.routingTable.NearestPeers(kb.ConvertKey(pmes.GetKey()), count)
	return closer
}

// betterPeerToQuery returns nearestPeersToQuery, but iff closer than self.
func (dht *IpfsDHT) betterPeersToQuery(pmes *pb.Message, p peer.ID, count int) []peer.ID {
	closer := dht.nearestPeersToQuery(pmes, count)

	// no node? nil
	if closer == nil {
		log.Warning("no closer peers to send:", p)
		return nil
	}

	filtered := make([]peer.ID, 0, len(closer))
	for _, clp := range closer {

		// == to self? thats bad
		if clp == dht.self {
			log.Warning("attempted to return self! this shouldn't happen...")
			return nil
		}
		// Dont send a peer back themselves
		if clp == p {
			continue
		}

		filtered = append(filtered, clp)
	}

	// ok seems like closer nodes
	return filtered
}

// Context return dht's context
func (dht *IpfsDHT) Context() context.Context {
	return dht.ctx
}

// Process return dht's process
func (dht *IpfsDHT) Process() goprocess.Process {
	return dht.proc
}

// Close calls Process Close
func (dht *IpfsDHT) Close() error {
	return dht.proc.Close()
}

func mkDsKey(s string) ds.Key {
	return ds.NewKey(base32.RawStdEncoding.EncodeToString([]byte(s)))
}
<|endoftext|>


<commit_before>/*
Command line tool for generating a Cloud Build yaml file based on versions.yaml.
*/
package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/GoogleCloudPlatform/runtimes-common/versioning/versions"
)

type cloudBuildOptions struct {
	// Whether to restrict to a particular set of Dockerfile directories.
	// If empty, all directories are used.
	Directories []string

	// Whether to run tests as part of the build.
	RunTests bool

	// Whether to require that image tags do not already exist in the repo.
	RequireNewTags bool

	// Whether to push to all declared tags
	FirstTagOnly bool

	// Optional timeout duration. If not specified, the Cloud Builder default timeout is used.
	TimeoutSeconds int

	// Optional parallel build. If specified, images can be build on bigger machines in parallel.
	EnableParallel bool

	// Forces parallel build. If specified, images are build on bigger machines in parallel. Overrides EnableParallel.
	ForceParallel bool
}

// TODO(huyhg): Replace "gcr.io/$PROJECT_ID/functional_test" with gcp-runtimes one.
const cloudBuildTemplateString = `steps:
{{- $parallel := .Parallel }}
{{- if .RequireNewTags }}
# Check if tags exist.
{{- range .Images }}
  - name: gcr.io/gcp-runtimes/check_if_tag_exists
    args:
      - 'python'
      - '/main.py'
      - '--image={{ . }}'
{{- end }}
{{- end }}

# Build images
{{- range .ImageBuilds }}
{{- if .Builder }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'build'
      - '--tag={{ .Tag }}'
      - '{{ .Directory }}'
{{- if $parallel }}
    waitFor: ['-']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- else }}
{{- if .BuilderImage }}
  - name: {{ .BuilderImage }}
    args: {{ .BuilderArgs }}
{{- if $parallel }}
    waitFor: ['image-{{ .BuilderImage }}']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- else }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'build'
      - '--tag={{ .Tag }}'
      - '{{ .Directory }}'
      - '{{ .Builder }}'
{{- if $parallel }}
    waitFor: ['-']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- end }}
{{- end }}
{{- end }}

{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range $testIndex, $test := $image.StructureTests }}
{{- if and (eq $imageIndex 0) (eq $testIndex 0) }}

# Run structure tests
{{- end}}
  - name: gcr.io/gcp-runtimes/structure_test
    args:
      - '--image'
      - '{{ $primary }}'
      - '--config'
      - '{{ $test }}'
{{- end }}
{{- end }}

{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range $testIndex, $test := $image.FunctionalTests }}
{{- if and (eq $imageIndex 0) (eq $testIndex 0) }}

# Run functional tests
{{- end }}
  - name: gcr.io/$PROJECT_ID/functional_test
    args:
      - '--verbose'
      - '--vars'
      - 'IMAGE={{ $primary }}'
      - '--vars'
      - 'UNIQUE={{ $imageIndex }}-{{ $testIndex }}'
      - '--test_spec'
      - '{{ $test }}'
{{- if $parallel }}
    waitFor: ['image-{{ $primary }}']
    id: 'test-{{ $primary }}-{{ $testIndex }}'
{{- end }}
{{- end }}

{{- end }}

# Add alias tags
{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range .Aliases }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'tag'
      - '{{ $primary }}'
      - '{{ . }}'
{{- if $parallel }}
    waitFor:
      - 'image-{{ $primary }}'
{{- range $testIndex, $test := $image.FunctionalTests }}
      - 'test-{{ $primary }}-{{ $testIndex }}'
{{- end }}
{{- end }}
{{- end }}
{{- end }}

images:
{{- range .AllImages }}
  - '{{ . }}'
{{- end }}

{{- if not (eq .TimeoutSeconds 0) }}

timeout: {{ .TimeoutSeconds }}s
{{- end }}

{{- if $parallel }}
options:
  machineType: 'N1_HIGHCPU_8'
{{- end }}
`

const testsDir = "tests"
const functionalTestsDir = "tests/functional_tests"
const structureTestsDir = "tests/structure_tests"
const testJsonSuffix = "_test.json"
const testYamlSuffix = "_test.yaml"
const workspacePrefix = "/workspace/"

type imageBuildTemplateData struct {
	Directory            string
	Tag                  string
	Aliases              []string
	StructureTests       []string
	FunctionalTests      []string
	Builder              bool
	BuilderImage         string
	BuilderArgs          []string
	ImageNameFromBuilder string
}

type cloudBuildTemplateData struct {
	RequireNewTags bool
	Parallel       bool
	ImageBuilds    []imageBuildTemplateData
	AllImages      []string
	TimeoutSeconds int
}

func shouldParallelize(options cloudBuildOptions, numberOfVersions int, numberOfTests int) bool {
	if options.ForceParallel {
		return true
	}
	if !options.EnableParallel {
		return false
	}
	return numberOfVersions > 1 || numberOfTests > 1
}

func newCloudBuildTemplateData(
	registry string, spec versions.Spec, options cloudBuildOptions) cloudBuildTemplateData {
	data := cloudBuildTemplateData{}
	data.RequireNewTags = options.RequireNewTags

	// Determine the set of directories to operate on.
	dirs := make(map[string]bool)
	if len(options.Directories) > 0 {
		for _, d := range options.Directories {
			dirs[d] = true
		}
	} else {
		for _, v := range spec.Versions {
			dirs[v.Dir] = true
		}
	}

	// Extract tests to run.
	var structureTests []string
	var functionalTests []string
	if options.RunTests {
		// Legacy structure tests reside in the root tests/ directory.
		structureTests = append(structureTests, readTests(testsDir)...)
		structureTests = append(structureTests, readTests(structureTestsDir)...)
		functionalTests = append(functionalTests, readTests(functionalTestsDir)...)
	}

	// Extract a list of full image names to build.
	for _, v := range spec.Versions {
		if !dirs[v.Dir] {
			continue
		}
		var images []string
		for _, t := range v.Tags {
			image := fmt.Sprintf("%v/%v:%v", registry, v.Repo, t)
			images = append(images, image)
			if options.FirstTagOnly {
				break
			}
		}
		// Ignore builder images from images list
		if !v.Builder {
			data.AllImages = append(data.AllImages, images...)
		}
		versionSTests, versionFTests := filterTests(structureTests, functionalTests, v)
		// Enforce to use ImageNameFromBuilder as reference to create tags
		if v.BuilderImage != "" {
			BuilderImageFull := fmt.Sprintf("%v/%v", registry, v.BuilderImage)
			data.ImageBuilds = append(
				data.ImageBuilds, imageBuildTemplateData{v.Dir, v.ImageNameFromBuilder, images, versionSTests, versionFTests, v.Builder, BuilderImageFull, v.BuilderArgs, v.ImageNameFromBuilder})
		} else {
			data.ImageBuilds = append(
				data.ImageBuilds, imageBuildTemplateData{v.Dir, images[0], images[1:], versionSTests, versionFTests, v.Builder, v.BuilderImage, v.BuilderArgs, v.ImageNameFromBuilder})
		}
	}

	data.TimeoutSeconds = options.TimeoutSeconds
	data.Parallel = shouldParallelize(options, len(spec.Versions), len(functionalTests))
	return data
}

func readTests(testsDir string) (tests []string) {
	if info, err := os.Stat(testsDir); err == nil && info.IsDir() {
		files, err := ioutil.ReadDir(testsDir)
		check(err)
		for _, f := range files {
			if f.IsDir() {
				continue
			}
			if strings.HasSuffix(f.Name(), testJsonSuffix) || strings.HasSuffix(f.Name(), testYamlSuffix) {
				tests = append(tests, workspacePrefix+fmt.Sprintf("%s/%s", testsDir, f.Name()))
			}
		}
	}
	return
}

func filterTests(structureTests []string, functionalTests []string, version versions.Version) (outStructureTests []string, outFunctionalTests []string) {
	included := make(map[string]bool, len(structureTests)+len(functionalTests))
	for _, test := range append(structureTests, functionalTests...) {
		included[test] = true
	}
	for _, excluded := range version.ExcludeTests {
		if !included[workspacePrefix+excluded] {
			log.Fatalf("No such test to exclude: %s", excluded)
		}
		included[workspacePrefix+excluded] = false
	}

	outStructureTests = make([]string, 0, len(structureTests))
	for _, test := range structureTests {
		if included[test] {
			outStructureTests = append(outStructureTests, test)
		}
	}
	outFunctionalTests = make([]string, 0, len(functionalTests))
	for _, test := range functionalTests {
		if included[test] {
			outFunctionalTests = append(outFunctionalTests, test)
		}
	}
	return
}

func renderCloudBuildConfig(
	registry string, spec versions.Spec, options cloudBuildOptions) string {
	data := newCloudBuildTemplateData(registry, spec, options)
	tmpl, _ := template.
		New("cloudBuildTemplate").
		Parse(cloudBuildTemplateString)
	var result bytes.Buffer
	tmpl.Execute(&result, data)
	return result.String()
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func main() {
	config := versions.LoadConfig("versions.yaml", "cloudbuild")
	registryPtr := config.StringOption("registry", "gcr.io/$PROJECT_ID", "Registry, e.g: 'gcr.io/my-project'")
	dirsPtr := config.StringOption("dirs", "", "Comma separated list of Dockerfile dirs to use.")
	testsPtr := config.BoolOption("tests", true, "Run tests.")
	newTagsPtr := config.BoolOption("new_tags", false, "Require that image tags do not already exist.")
	firstTagOnly := config.BoolOption("first_tag", false, "Build only the first per version.")
	timeoutPtr := config.IntOption("timeout", 0, "Timeout in seconds. If not set, the default Cloud Build timeout is used.")
	enableParallel := config.BoolOption("enable_parallel", false, "Enable parallel build and bigger VM")
	forceParallel := config.BoolOption("force_parallel", false, "Force parallel build and bigger VM")
	config.Parse()

	if *registryPtr == "" {
		log.Fatalf("--registry flag is required")
	}

	if strings.Contains(*registryPtr, ":") {
		*registryPtr = strings.Replace(*registryPtr, ":", "/", 1)
	}

	var dirs []string
	if *dirsPtr != "" {
		dirs = strings.Split(*dirsPtr, ",")
	}
	spec := versions.LoadVersions("versions.yaml")
	options := cloudBuildOptions{dirs, *testsPtr, *newTagsPtr, *firstTagOnly, *timeoutPtr, *enableParallel, *forceParallel}
	result := renderCloudBuildConfig(*registryPtr, spec, options)
	fmt.Println(result)
}
<commit_msg>fix template (#802)<commit_after>/*
Command line tool for generating a Cloud Build yaml file based on versions.yaml.
*/
package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/GoogleCloudPlatform/runtimes-common/versioning/versions"
)

type cloudBuildOptions struct {
	// Whether to restrict to a particular set of Dockerfile directories.
	// If empty, all directories are used.
	Directories []string

	// Whether to run tests as part of the build.
	RunTests bool

	// Whether to require that image tags do not already exist in the repo.
	RequireNewTags bool

	// Whether to push to all declared tags
	FirstTagOnly bool

	// Optional timeout duration. If not specified, the Cloud Builder default timeout is used.
	TimeoutSeconds int

	// Optional parallel build. If specified, images can be build on bigger machines in parallel.
	EnableParallel bool

	// Forces parallel build. If specified, images are build on bigger machines in parallel. Overrides EnableParallel.
	ForceParallel bool
}

// TODO(huyhg): Replace "gcr.io/$PROJECT_ID/functional_test" with gcp-runtimes one.
const cloudBuildTemplateString = `steps:
{{- $parallel := .Parallel }}
{{- if .RequireNewTags }}
# Check if tags exist.
{{- range .Images }}
  - name: gcr.io/gcp-runtimes/check_if_tag_exists
    args:
      - 'python'
      - '/main.py'
      - '--image={{ . }}'
{{- end }}
{{- end }}

# Build images
{{- range .ImageBuilds }}
{{- if .Builder }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'build'
      - '--tag={{ .Tag }}'
      - '{{ .Directory }}'
{{- if $parallel }}
    waitFor: ['-']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- else }}
{{- if .BuilderImage }}
  - name: {{ .BuilderImage }}
    args: {{ .BuilderArgs }}
{{- if $parallel }}
    waitFor: ['image-{{ .BuilderImage }}']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- else }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'build'
      - '--tag={{ .Tag }}'
      - '{{ .Directory }}'
{{- if $parallel }}
    waitFor: ['-']
    id: 'image-{{ .Tag }}'
{{- end }}
{{- end }}
{{- end }}
{{- end }}

{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range $testIndex, $test := $image.StructureTests }}
{{- if and (eq $imageIndex 0) (eq $testIndex 0) }}

# Run structure tests
{{- end}}
  - name: gcr.io/gcp-runtimes/structure_test
    args:
      - '--image'
      - '{{ $primary }}'
      - '--config'
      - '{{ $test }}'
{{- end }}
{{- end }}

{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range $testIndex, $test := $image.FunctionalTests }}
{{- if and (eq $imageIndex 0) (eq $testIndex 0) }}

# Run functional tests
{{- end }}
  - name: gcr.io/$PROJECT_ID/functional_test
    args:
      - '--verbose'
      - '--vars'
      - 'IMAGE={{ $primary }}'
      - '--vars'
      - 'UNIQUE={{ $imageIndex }}-{{ $testIndex }}'
      - '--test_spec'
      - '{{ $test }}'
{{- if $parallel }}
    waitFor: ['image-{{ $primary }}']
    id: 'test-{{ $primary }}-{{ $testIndex }}'
{{- end }}
{{- end }}

{{- end }}

# Add alias tags
{{- range $imageIndex, $image := .ImageBuilds }}
{{- $primary := $image.Tag }}
{{- range .Aliases }}
  - name: gcr.io/cloud-builders/docker
    args:
      - 'tag'
      - '{{ $primary }}'
      - '{{ . }}'
{{- if $parallel }}
    waitFor:
      - 'image-{{ $primary }}'
{{- range $testIndex, $test := $image.FunctionalTests }}
      - 'test-{{ $primary }}-{{ $testIndex }}'
{{- end }}
{{- end }}
{{- end }}
{{- end }}

images:
{{- range .AllImages }}
  - '{{ . }}'
{{- end }}

{{- if not (eq .TimeoutSeconds 0) }}

timeout: {{ .TimeoutSeconds }}s
{{- end }}

{{- if $parallel }}
options:
  machineType: 'N1_HIGHCPU_8'
{{- end }}
`

const testsDir = "tests"
const functionalTestsDir = "tests/functional_tests"
const structureTestsDir = "tests/structure_tests"
const testJsonSuffix = "_test.json"
const testYamlSuffix = "_test.yaml"
const workspacePrefix = "/workspace/"

type imageBuildTemplateData struct {
	Directory            string
	Tag                  string
	Aliases              []string
	StructureTests       []string
	FunctionalTests      []string
	Builder              bool
	BuilderImage         string
	BuilderArgs          []string
	ImageNameFromBuilder string
}

type cloudBuildTemplateData struct {
	RequireNewTags bool
	Parallel       bool
	ImageBuilds    []imageBuildTemplateData
	AllImages      []string
	TimeoutSeconds int
}

func shouldParallelize(options cloudBuildOptions, numberOfVersions int, numberOfTests int) bool {
	if options.ForceParallel {
		return true
	}
	if !options.EnableParallel {
		return false
	}
	return numberOfVersions > 1 || numberOfTests > 1
}

func newCloudBuildTemplateData(
	registry string, spec versions.Spec, options cloudBuildOptions) cloudBuildTemplateData {
	data := cloudBuildTemplateData{}
	data.RequireNewTags = options.RequireNewTags

	// Determine the set of directories to operate on.
	dirs := make(map[string]bool)
	if len(options.Directories) > 0 {
		for _, d := range options.Directories {
			dirs[d] = true
		}
	} else {
		for _, v := range spec.Versions {
			dirs[v.Dir] = true
		}
	}

	// Extract tests to run.
	var structureTests []string
	var functionalTests []string
	if options.RunTests {
		// Legacy structure tests reside in the root tests/ directory.
		structureTests = append(structureTests, readTests(testsDir)...)
		structureTests = append(structureTests, readTests(structureTestsDir)...)
		functionalTests = append(functionalTests, readTests(functionalTestsDir)...)
	}

	// Extract a list of full image names to build.
	for _, v := range spec.Versions {
		if !dirs[v.Dir] {
			continue
		}
		var images []string
		for _, t := range v.Tags {
			image := fmt.Sprintf("%v/%v:%v", registry, v.Repo, t)
			images = append(images, image)
			if options.FirstTagOnly {
				break
			}
		}
		// Ignore builder images from images list
		if !v.Builder {
			data.AllImages = append(data.AllImages, images...)
		}
		versionSTests, versionFTests := filterTests(structureTests, functionalTests, v)
		// Enforce to use ImageNameFromBuilder as reference to create tags
		if v.BuilderImage != "" {
			BuilderImageFull := fmt.Sprintf("%v/%v", registry, v.BuilderImage)
			data.ImageBuilds = append(
				data.ImageBuilds, imageBuildTemplateData{v.Dir, v.ImageNameFromBuilder, images, versionSTests, versionFTests, v.Builder, BuilderImageFull, v.BuilderArgs, v.ImageNameFromBuilder})
		} else {
			data.ImageBuilds = append(
				data.ImageBuilds, imageBuildTemplateData{v.Dir, images[0], images[1:], versionSTests, versionFTests, v.Builder, v.BuilderImage, v.BuilderArgs, v.ImageNameFromBuilder})
		}
	}

	data.TimeoutSeconds = options.TimeoutSeconds
	data.Parallel = shouldParallelize(options, len(spec.Versions), len(functionalTests))
	return data
}

func readTests(testsDir string) (tests []string) {
	if info, err := os.Stat(testsDir); err == nil && info.IsDir() {
		files, err := ioutil.ReadDir(testsDir)
		check(err)
		for _, f := range files {
			if f.IsDir() {
				continue
			}
			if strings.HasSuffix(f.Name(), testJsonSuffix) || strings.HasSuffix(f.Name(), testYamlSuffix) {
				tests = append(tests, workspacePrefix+fmt.Sprintf("%s/%s", testsDir, f.Name()))
			}
		}
	}
	return
}

func filterTests(structureTests []string, functionalTests []string, version versions.Version) (outStructureTests []string, outFunctionalTests []string) {
	included := make(map[string]bool, len(structureTests)+len(functionalTests))
	for _, test := range append(structureTests, functionalTests...) {
		included[test] = true
	}
	for _, excluded := range version.ExcludeTests {
		if !included[workspacePrefix+excluded] {
			log.Fatalf("No such test to exclude: %s", excluded)
		}
		included[workspacePrefix+excluded] = false
	}

	outStructureTests = make([]string, 0, len(structureTests))
	for _, test := range structureTests {
		if included[test] {
			outStructureTests = append(outStructureTests, test)
		}
	}
	outFunctionalTests = make([]string, 0, len(functionalTests))
	for _, test := range functionalTests {
		if included[test] {
			outFunctionalTests = append(outFunctionalTests, test)
		}
	}
	return
}

func renderCloudBuildConfig(
	registry string, spec versions.Spec, options cloudBuildOptions) string {
	data := newCloudBuildTemplateData(registry, spec, options)
	tmpl, _ := template.
		New("cloudBuildTemplate").
		Parse(cloudBuildTemplateString)
	var result bytes.Buffer
	tmpl.Execute(&result, data)
	return result.String()
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}

func main() {
	config := versions.LoadConfig("versions.yaml", "cloudbuild")
	registryPtr := config.StringOption("registry", "gcr.io/$PROJECT_ID", "Registry, e.g: 'gcr.io/my-project'")
	dirsPtr := config.StringOption("dirs", "", "Comma separated list of Dockerfile dirs to use.")
	testsPtr := config.BoolOption("tests", true, "Run tests.")
	newTagsPtr := config.BoolOption("new_tags", false, "Require that image tags do not already exist.")
	firstTagOnly := config.BoolOption("first_tag", false, "Build only the first per version.")
	timeoutPtr := config.IntOption("timeout", 0, "Timeout in seconds. If not set, the default Cloud Build timeout is used.")
	enableParallel := config.BoolOption("enable_parallel", false, "Enable parallel build and bigger VM")
	forceParallel := config.BoolOption("force_parallel", false, "Force parallel build and bigger VM")
	config.Parse()

	if *registryPtr == "" {
		log.Fatalf("--registry flag is required")
	}

	if strings.Contains(*registryPtr, ":") {
		*registryPtr = strings.Replace(*registryPtr, ":", "/", 1)
	}

	var dirs []string
	if *dirsPtr != "" {
		dirs = strings.Split(*dirsPtr, ",")
	}
	spec := versions.LoadVersions("versions.yaml")
	options := cloudBuildOptions{dirs, *testsPtr, *newTagsPtr, *firstTagOnly, *timeoutPtr, *enableParallel, *forceParallel}
	result := renderCloudBuildConfig(*registryPtr, spec, options)
	fmt.Println(result)
}
<|endoftext|>


<commit_before>package protocol

import "testing"

func TestBrokerEquals(t *testing.T) {
	var b1, b2 *Broker

	b1 = nil
	b2 = nil

	if !b1.Equals(b2) {
		t.Error("Two nil brokers didn't compare equal.")
	}

	b1 = NewBroker("abc", 123)

	if b1.Equals(b2) {
		t.Error("Non-nil and nil brokers compared equal.")
	}
	if b2.Equals(b1) {
		t.Error("Nil and non-nil brokers compared equal.")
	}

	b2 = NewBroker("abc", 1234)
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different ports compared equal.")
	}

	b2 = NewBroker("abcd", 123)
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different hosts compared equal.")
	}

	b2 = NewBroker("abc", 123)
	b2.id = -2
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different ids compared equal.")
	}

	b2.id = -1
	if !b1.Equals(b2) || !b2.Equals(b1) {
		t.Error("Similar brokers did not compare equal.")
	}
}
<commit_msg>Basic broker connect/close test, including framework for dummy server<commit_after>package protocol

import (
	"encoding/binary"
	"io"
	"net"
	"strconv"
	"testing"
)

func fakeTCPServer(t *testing.T, responses [][]byte, done chan<- bool) (int32, error) {
	ln, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		return 0, err
	}
	_, portStr, err := net.SplitHostPort(ln.Addr().String())
	if err != nil {
		return 0, err
	}
	tmp, err := strconv.ParseInt(portStr, 10, 32)
	if err != nil {
		return 0, err
	}
	port := int32(tmp)
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			t.Error(err)
			conn.Close()
			ln.Close()
			done<- true
			return
		}
		for _, response := range responses {
			header := make([]byte, 4)
			_, err := io.ReadFull(conn, header)
			if err != nil {
				t.Error(err)
				conn.Close()
				ln.Close()
				done<- true
				return
			}
			body := make([]byte, binary.BigEndian.Uint32(header))
			_, err = io.ReadFull(conn, body)
			if err != nil {
				t.Error(err)
				conn.Close()
				ln.Close()
				done<- true
				return
			}
			_, err = conn.Write(response)
			if err != nil {
				t.Error(err)
				conn.Close()
				ln.Close()
				done<- true
				return
			}
		}
		err = conn.Close()
		if err != nil {
			t.Error(err)
			ln.Close()
			done<- true
			return
		}
		err = ln.Close()
		if err != nil {
			t.Error(err)
			done<- true
			return
		}
		done<- true
	}()
	return port, nil
}

func TestBrokerEquals(t *testing.T) {
	var b1, b2 *Broker

	b1 = nil
	b2 = nil

	if !b1.Equals(b2) {
		t.Error("Two nil brokers didn't compare equal.")
	}

	b1 = NewBroker("abc", 123)

	if b1.Equals(b2) {
		t.Error("Non-nil and nil brokers compared equal.")
	}
	if b2.Equals(b1) {
		t.Error("Nil and non-nil brokers compared equal.")
	}

	b2 = NewBroker("abc", 1234)
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different ports compared equal.")
	}

	b2 = NewBroker("abcd", 123)
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different hosts compared equal.")
	}

	b2 = NewBroker("abc", 123)
	b2.id = -2
	if b1.Equals(b2) || b2.Equals(b1) {
		t.Error("Brokers with different ids compared equal.")
	}

	b2.id = -1
	if !b1.Equals(b2) || !b2.Equals(b1) {
		t.Error("Similar brokers did not compare equal.")
	}
}

func TestBrokerID(t *testing.T) {

	broker := NewBroker("abc", 123)

	if broker.ID() != -1 {
		t.Error("New broker didn't have an ID of -1.")
	}

	broker.id = 34
	if broker.ID() != 34 {
		t.Error("Manually setting broker ID did not take effect.")
	}
}

func TestBrokerConnectClose(t *testing.T) {
	done := make(chan bool)
	port, err := fakeTCPServer(t, [][]byte{}, done)
	if err != nil {
		t.Error(err)
		return
	}
	broker := NewBroker("localhost", port)
	err = broker.Connect()
	if err != nil {
		t.Error(err)
		return
	}
	err = broker.Close()
	if err != nil {
		t.Error(err)
		return
	}
	<-done
}
<|endoftext|>


<commit_before>package identify

import (
	"fmt"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	ggio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/gogoprotobuf/io"
	semver "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/coreos/go-semver/semver"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"

	host "github.com/jbenet/go-ipfs/p2p/host"
	inet "github.com/jbenet/go-ipfs/p2p/net"
	protocol "github.com/jbenet/go-ipfs/p2p/protocol"
	pb "github.com/jbenet/go-ipfs/p2p/protocol/identify/pb"
	config "github.com/jbenet/go-ipfs/repo/config"
	eventlog "github.com/jbenet/go-ipfs/thirdparty/eventlog"
)

var log = eventlog.Logger("net/identify")

// ID is the protocol.ID of the Identify Service.
const ID protocol.ID = "/ipfs/identify"

// IpfsVersion holds the current protocol version for a client running this code
var IpfsVersion *semver.Version
var ClientVersion = "go-ipfs/" + config.CurrentVersionNumber

func init() {
	var err error
	IpfsVersion, err = semver.NewVersion("0.0.1")
	if err != nil {
		panic(fmt.Errorf("invalid protocol version: %v", err))
	}
}

// IDService is a structure that implements ProtocolIdentify.
// It is a trivial service that gives the other peer some
// useful information about the local peer. A sort of hello.
//
// The IDService sends:
//  * Our IPFS Protocol Version
//  * Our IPFS Agent Version
//  * Our public Listen Addresses
type IDService struct {
	Host host.Host

	// connections undergoing identification
	// for wait purposes
	currid map[inet.Conn]chan struct{}
	currmu sync.RWMutex
}

func NewIDService(h host.Host) *IDService {
	s := &IDService{
		Host:   h,
		currid: make(map[inet.Conn]chan struct{}),
	}
	h.SetStreamHandler(ID, s.RequestHandler)
	return s
}

func (ids *IDService) IdentifyConn(c inet.Conn) {
	ids.currmu.Lock()
	if wait, found := ids.currid[c]; found {
		ids.currmu.Unlock()
		log.Debugf("IdentifyConn called twice on: %s", c)
		<-wait // already identifying it. wait for it.
		return
	}
	ids.currid[c] = make(chan struct{})
	ids.currmu.Unlock()

	s, err := c.NewStream()
	if err != nil {
		log.Error("error opening initial stream for %s", ID)
		log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
	} else {

		// ok give the response to our handler.
		if err := protocol.WriteHeader(s, ID); err != nil {
			log.Error("error writing stream header for %s", ID)
			log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
		}
		ids.ResponseHandler(s)
	}

	ids.currmu.Lock()
	ch, found := ids.currid[c]
	delete(ids.currid, c)
	ids.currmu.Unlock()

	if !found {
		log.Errorf("IdentifyConn failed to find channel (programmer error) for %s", c)
		return
	}

	close(ch) // release everyone waiting.
}

func (ids *IDService) RequestHandler(s inet.Stream) {
	defer s.Close()
	c := s.Conn()

	w := ggio.NewDelimitedWriter(s)
	mes := pb.Identify{}
	ids.populateMessage(&mes, s.Conn())
	w.WriteMsg(&mes)

	log.Debugf("%s sent message to %s %s", ID,
		c.RemotePeer(), c.RemoteMultiaddr())
}

func (ids *IDService) ResponseHandler(s inet.Stream) {
	defer s.Close()
	c := s.Conn()

	r := ggio.NewDelimitedReader(s, 2048)
	mes := pb.Identify{}
	if err := r.ReadMsg(&mes); err != nil {
		log.Errorf("%s error receiving message from %s %s", ID,
			c.RemotePeer(), c.RemoteMultiaddr())
		return
	}
	ids.consumeMessage(&mes, c)

	log.Debugf("%s received message from %s %s", ID,
		c.RemotePeer(), c.RemoteMultiaddr())
}

func (ids *IDService) populateMessage(mes *pb.Identify, c inet.Conn) {

	// set protocols this node is currently handling
	protos := ids.Host.Mux().Protocols()
	mes.Protocols = make([]string, len(protos))
	for i, p := range protos {
		mes.Protocols[i] = string(p)
	}

	// observed address so other side is informed of their
	// "public" address, at least in relation to us.
	mes.ObservedAddr = c.RemoteMultiaddr().Bytes()

	// set listen addrs
	laddrs, err := ids.Host.Network().InterfaceListenAddresses()
	if err != nil {
		log.Error(err)
	} else {
		mes.ListenAddrs = make([][]byte, len(laddrs))
		for i, addr := range laddrs {
			mes.ListenAddrs[i] = addr.Bytes()
		}
		log.Debugf("%s sent listen addrs to %s: %s", c.LocalPeer(), c.RemotePeer(), laddrs)
	}

	// set protocol versions
	s := IpfsVersion.String()
	mes.ProtocolVersion = &s
	mes.AgentVersion = &ClientVersion
}

func (ids *IDService) consumeMessage(mes *pb.Identify, c inet.Conn) {
	p := c.RemotePeer()

	// mes.Protocols
	// mes.ObservedAddr

	// mes.ListenAddrs
	laddrs := mes.GetListenAddrs()
	lmaddrs := make([]ma.Multiaddr, 0, len(laddrs))
	for _, addr := range laddrs {
		maddr, err := ma.NewMultiaddrBytes(addr)
		if err != nil {
			log.Errorf("%s failed to parse multiaddr from %s %s", ID,
				p, c.RemoteMultiaddr())
			continue
		}
		lmaddrs = append(lmaddrs, maddr)
	}

	// update our peerstore with the addresses.
	ids.Host.Peerstore().AddAddresses(p, lmaddrs)
	log.Debugf("%s received listen addrs for %s: %s", c.LocalPeer(), c.RemotePeer(), lmaddrs)

	// get protocol versions
	pv := *mes.ProtocolVersion
	av := *mes.AgentVersion
	ids.Host.Peerstore().Put(p, "ProtocolVersion", pv)
	ids.Host.Peerstore().Put(p, "AgentVersion", av)
}

// IdentifyWait returns a channel which will be closed once
// "ProtocolIdentify" (handshake3) finishes on given conn.
// This happens async so the connection can start to be used
// even if handshake3 knowledge is not necesary.
// Users **MUST** call IdentifyWait _after_ IdentifyConn
func (ids *IDService) IdentifyWait(c inet.Conn) <-chan struct{} {
	ids.currmu.Lock()
	ch, found := ids.currid[c]
	ids.currmu.Unlock()
	if found {
		return ch
	}

	// if not found, it means we are already done identifying it, or
	// haven't even started. either way, return a new channel closed.
	ch = make(chan struct{})
	close(ch)
	return ch
}
<commit_msg>p2p/proto/identify: use observed listen addrs<commit_after>package identify

import (
	"fmt"
	"sync"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	ggio "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/gogoprotobuf/io"
	semver "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/coreos/go-semver/semver"
	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"

	host "github.com/jbenet/go-ipfs/p2p/host"
	inet "github.com/jbenet/go-ipfs/p2p/net"
	protocol "github.com/jbenet/go-ipfs/p2p/protocol"
	pb "github.com/jbenet/go-ipfs/p2p/protocol/identify/pb"
	config "github.com/jbenet/go-ipfs/repo/config"
	eventlog "github.com/jbenet/go-ipfs/thirdparty/eventlog"
)

var log = eventlog.Logger("net/identify")

// ID is the protocol.ID of the Identify Service.
const ID protocol.ID = "/ipfs/identify"

// IpfsVersion holds the current protocol version for a client running this code
var IpfsVersion *semver.Version
var ClientVersion = "go-ipfs/" + config.CurrentVersionNumber

func init() {
	var err error
	IpfsVersion, err = semver.NewVersion("0.0.1")
	if err != nil {
		panic(fmt.Errorf("invalid protocol version: %v", err))
	}
}

// IDService is a structure that implements ProtocolIdentify.
// It is a trivial service that gives the other peer some
// useful information about the local peer. A sort of hello.
//
// The IDService sends:
//  * Our IPFS Protocol Version
//  * Our IPFS Agent Version
//  * Our public Listen Addresses
type IDService struct {
	Host host.Host

	// connections undergoing identification
	// for wait purposes
	currid map[inet.Conn]chan struct{}
	currmu sync.RWMutex
}

func NewIDService(h host.Host) *IDService {
	s := &IDService{
		Host:   h,
		currid: make(map[inet.Conn]chan struct{}),
	}
	h.SetStreamHandler(ID, s.RequestHandler)
	return s
}

func (ids *IDService) IdentifyConn(c inet.Conn) {
	ids.currmu.Lock()
	if wait, found := ids.currid[c]; found {
		ids.currmu.Unlock()
		log.Debugf("IdentifyConn called twice on: %s", c)
		<-wait // already identifying it. wait for it.
		return
	}
	ids.currid[c] = make(chan struct{})
	ids.currmu.Unlock()

	s, err := c.NewStream()
	if err != nil {
		log.Error("error opening initial stream for %s", ID)
		log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
	} else {

		// ok give the response to our handler.
		if err := protocol.WriteHeader(s, ID); err != nil {
			log.Error("error writing stream header for %s", ID)
			log.Event(context.TODO(), "IdentifyOpenFailed", c.RemotePeer())
		}
		ids.ResponseHandler(s)
	}

	ids.currmu.Lock()
	ch, found := ids.currid[c]
	delete(ids.currid, c)
	ids.currmu.Unlock()

	if !found {
		log.Errorf("IdentifyConn failed to find channel (programmer error) for %s", c)
		return
	}

	close(ch) // release everyone waiting.
}

func (ids *IDService) RequestHandler(s inet.Stream) {
	defer s.Close()
	c := s.Conn()

	w := ggio.NewDelimitedWriter(s)
	mes := pb.Identify{}
	ids.populateMessage(&mes, s.Conn())
	w.WriteMsg(&mes)

	log.Debugf("%s sent message to %s %s", ID,
		c.RemotePeer(), c.RemoteMultiaddr())
}

func (ids *IDService) ResponseHandler(s inet.Stream) {
	defer s.Close()
	c := s.Conn()

	r := ggio.NewDelimitedReader(s, 2048)
	mes := pb.Identify{}
	if err := r.ReadMsg(&mes); err != nil {
		log.Errorf("%s error receiving message from %s %s", ID,
			c.RemotePeer(), c.RemoteMultiaddr())
		return
	}
	ids.consumeMessage(&mes, c)

	log.Debugf("%s received message from %s %s", ID,
		c.RemotePeer(), c.RemoteMultiaddr())
}

func (ids *IDService) populateMessage(mes *pb.Identify, c inet.Conn) {

	// set protocols this node is currently handling
	protos := ids.Host.Mux().Protocols()
	mes.Protocols = make([]string, len(protos))
	for i, p := range protos {
		mes.Protocols[i] = string(p)
	}

	// observed address so other side is informed of their
	// "public" address, at least in relation to us.
	mes.ObservedAddr = c.RemoteMultiaddr().Bytes()

	// set listen addrs
	laddrs, err := ids.Host.Network().InterfaceListenAddresses()
	if err != nil {
		log.Error(err)
	} else {
		mes.ListenAddrs = make([][]byte, len(laddrs))
		for i, addr := range laddrs {
			mes.ListenAddrs[i] = addr.Bytes()
		}
		log.Debugf("%s sent listen addrs to %s: %s", c.LocalPeer(), c.RemotePeer(), laddrs)
	}

	// set protocol versions
	s := IpfsVersion.String()
	mes.ProtocolVersion = &s
	mes.AgentVersion = &ClientVersion
}

func (ids *IDService) consumeMessage(mes *pb.Identify, c inet.Conn) {
	p := c.RemotePeer()

	// mes.Protocols

	// mes.ObservedAddr
	ids.consumeObservedAddress(mes.GetObservedAddr(), c)

	// mes.ListenAddrs
	laddrs := mes.GetListenAddrs()
	lmaddrs := make([]ma.Multiaddr, 0, len(laddrs))
	for _, addr := range laddrs {
		maddr, err := ma.NewMultiaddrBytes(addr)
		if err != nil {
			log.Errorf("%s failed to parse multiaddr from %s %s", ID,
				p, c.RemoteMultiaddr())
			continue
		}
		lmaddrs = append(lmaddrs, maddr)
	}

	// update our peerstore with the addresses.
	ids.Host.Peerstore().AddAddresses(p, lmaddrs)
	log.Debugf("%s received listen addrs for %s: %s", c.LocalPeer(), c.RemotePeer(), lmaddrs)

	// get protocol versions
	pv := *mes.ProtocolVersion
	av := *mes.AgentVersion
	ids.Host.Peerstore().Put(p, "ProtocolVersion", pv)
	ids.Host.Peerstore().Put(p, "AgentVersion", av)
}

// IdentifyWait returns a channel which will be closed once
// "ProtocolIdentify" (handshake3) finishes on given conn.
// This happens async so the connection can start to be used
// even if handshake3 knowledge is not necesary.
// Users **MUST** call IdentifyWait _after_ IdentifyConn
func (ids *IDService) IdentifyWait(c inet.Conn) <-chan struct{} {
	ids.currmu.Lock()
	ch, found := ids.currid[c]
	ids.currmu.Unlock()
	if found {
		return ch
	}

	// if not found, it means we are already done identifying it, or
	// haven't even started. either way, return a new channel closed.
	ch = make(chan struct{})
	close(ch)
	return ch
}

func (ids *IDService) consumeObservedAddress(observed []byte, c inet.Conn) {
	if observed == nil {
		return
	}

	maddr, err := ma.NewMultiaddrBytes(observed)
	if err != nil {
		log.Debugf("error parsing received observed addr for %s: %s", c, err)
		return
	}

	// we should only use ObservedAddr when our connection's LocalAddr is one
	// of our ListenAddrs. If we Dial out using an ephemeral addr, knowing that
	// address's external mapping is not very useful because the port will not be
	// the same as the listen addr.
	ifaceaddrs, err := ids.Host.Network().InterfaceListenAddresses()
	if err != nil {
		log.Infof("failed to get interface listen addrs", err)
		return
	}

	log.Debugf("identify identifying observed multiaddr: %s %s", c.LocalMultiaddr(), ifaceaddrs)
	if !addrInAddrs(c.LocalMultiaddr(), ifaceaddrs) {
		// not in our list
		return
	}

	// ok! we have the observed version of one of our ListenAddresses!
	log.Debugf("added own observed listen addr: %s --> %s", c.LocalMultiaddr(), maddr)
	ids.Host.Peerstore().AddAddress(ids.Host.ID(), maddr)
}

func addrInAddrs(a ma.Multiaddr, as []ma.Multiaddr) bool {
	for _, b := range as {
		if a.Equal(b) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package protocol

import (
	"testing"
	"time"
)

var script = []TimedAction{
	{
		Action{
			Position: 5,
			Speed:    50,
		},
		time.Millisecond * 50,
	},
	{
		Action{
			Position: 50,
			Speed:    40,
		},
		time.Millisecond * 100,
	},
	{
		Action{
			Position: 90,
			Speed:    90,
		},
		time.Millisecond * 150,
	},
	{
		Action{
			Position: 30,
			Speed:    30,
		},
		time.Millisecond * 200,
	},
}

func TestPlay(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	var eventCount int
	starttime := time.Now()
	for a := range p.Play() {
		eventtime := time.Now().Sub(starttime)
		t.Logf("Action: %s: %d,%d", eventtime, a.Position, a.Speed)
		eventCount++
	}
	playTime := time.Now().Sub(starttime)

	if eventCount != len(script) {
		t.Errorf("not all actions were played, want %d, got %d",
			len(script), eventCount)
	}
	want := script[len(script)-1].Time
	if playTime.Nanoseconds()/1e6 != want.Nanoseconds()/1e6 {
		t.Errorf("script was not played back at correct speed")
	}
}

func TestPauseResume(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	pauseTime := time.Millisecond * 100

	go func() {
		<-time.After(time.Millisecond * 75)
		if err := p.Pause(); err != nil {
			t.Error(err)
		}
		<-time.After(pauseTime)
		if err := p.Resume(); err != nil {
			t.Error(err)
		}
	}()

	var eventCount int
	starttime := time.Now()
	for range p.Play() {
		eventCount++
	}
	playTime := time.Now().Sub(starttime)

	if eventCount != len(script) {
		t.Errorf("not all actions were played, want %d, got %d",
			len(script), eventCount)
	}
	want := script[len(script)-1].Time + pauseTime
	if playTime.Nanoseconds()/1e6 != want.Nanoseconds()/1e6 {
		t.Errorf("script was not played back at correct speed")
	}
}

func TestStop(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	stopTime := time.Millisecond * 75

	go func() {
		<-time.After(stopTime)
		if err := p.Stop(); err != nil {
			t.Error(err)
		}
	}()

	starttime := time.Now()
	for range p.Play() {
		// pass
	}
	playTime := time.Now().Sub(starttime)

	if playTime.Nanoseconds()/1e6 != stopTime.Nanoseconds()/1e6 {
		t.Errorf("script was not stopped at the right time, want: %s, got: %s",
			stopTime, playTime)
	}
}

func TestSkip(t *testing.T) {
	cases := []struct {
		Name string
		At   time.Duration
		To   time.Duration
	}{
		{
			Name: "Forward",
			At:   time.Millisecond * 50,
			To:   time.Millisecond * 150,
		},
		{
			Name: "Back",
			At:   time.Millisecond * 50,
			To:   time.Millisecond * 150,
		},
		{
			Name: "Soon",
			At:   time.Millisecond * 0,
			To:   time.Millisecond * 150,
		},
	}

	for _, c := range cases {
		p := NewTimedActionsPlayer()
		p.Script = script

		go func() {
			<-time.After(c.At)
			if err := p.Skip(c.To); err != nil {
				t.Error(err)
			}
		}()

		starttime := time.Now()
		for range p.Play() {
			// pass
		}
		playTime := time.Now().Sub(starttime)
		want := script[len(script)-1].Time - (c.To - c.At)

		if playTime.Nanoseconds()/1e6 != want.Nanoseconds()/1e6 {
			t.Errorf("%s: player did not skip correctly, want: %s, got: %s",
				c.Name, want, playTime)
		}
	}
}

func TestLimits(t *testing.T) {
	type TestCase struct {
		Low, High  int
		Slow, Fast int
	}

	cases := []TestCase{
		{0, 100, 0, 0},
		{10, 90, 0, 0},
		{30, 50, 0, 0},
		{50, 90, 0, 0},
		{60, 80, 0, 0},
		{0, 0, 20, 80},
		{0, 0, 30, 60},
		{0, 0, 50, 90},
		{0, 0, 10, 20},
		{0, 0, 0, 50},
		{0, 100, 20, 80},
		{10, 90, 30, 60},
		{30, 50, 50, 90},
		{50, 90, 10, 20},
		{60, 80, 0, 50},
	}
	for i, c := range cases {
		p := NewTimedActionsPlayer()
		p.Script = script
		if c.Low != 0 && c.High != 0 {
			p.LimitPosition(c.Low, c.High)
		}
		if c.Slow != 0 && c.Fast != 0 {
			p.LimitSpeed(c.Slow, c.Fast)
		}

		lowest, highest := c.Low, c.High
		slowest, fastest := c.Slow, c.Fast

		for a := range p.Play() {
			if a.Position < lowest {
				lowest = a.Position
			} else if a.Position > highest {
				highest = a.Position
			}
			if a.Position < slowest {
				slowest = a.Speed
			} else if a.Position > highest {
				fastest = a.Speed
			}
		}

		if c.Low != 0 && c.High != 0 {
			if lowest < c.Low {
				t.Errorf("case %d: went lower than allowed, %d < %d", i, lowest, c.Low)
			}
			if highest > c.High {
				t.Errorf("case %d: went higher than allowed, %d > %d", i, highest, c.High)
			}
		}

		if c.Slow != 0 && c.Fast != 0 {
			if slowest < c.Slow {
				t.Errorf("case %d: went slower than allowed, %d < %d", i, slowest, c.Slow)
			}
			if fastest > c.Fast {
				t.Errorf("case %d: went higher than allowed, %d > %d", i, fastest, c.Fast)
			}
		}

	}
}

func TestDump(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	actions, err := p.Dump()
	if err != nil {
		t.Error(err)
	}
	if len(script) != len(actions) {
		t.Errorf("dump did not return script actions: want %d, got %d",
			len(script), len(actions))
	}
}
<commit_msg>Use a bigger tolerance when doing timing test<commit_after>package protocol

import (
	"testing"
	"time"
)

var script = []TimedAction{
	{
		Action{
			Position: 5,
			Speed:    50,
		},
		time.Millisecond * 50,
	},
	{
		Action{
			Position: 50,
			Speed:    40,
		},
		time.Millisecond * 100,
	},
	{
		Action{
			Position: 90,
			Speed:    90,
		},
		time.Millisecond * 150,
	},
	{
		Action{
			Position: 30,
			Speed:    30,
		},
		time.Millisecond * 200,
	},
}

type timeTolerance time.Duration

func (p timeTolerance) roughlyEqual(a time.Duration, b time.Duration) bool {
	if a > b+time.Duration(p) {
		return false
	}
	if a < b-time.Duration(p) {
		return false
	}
	return true
}

var defaultTimeTolerance = timeTolerance(time.Millisecond * 10)

func TestPlay(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	var eventCount int
	starttime := time.Now()
	for a := range p.Play() {
		eventtime := time.Now().Sub(starttime)
		t.Logf("Action: %s: %d,%d", eventtime, a.Position, a.Speed)
		eventCount++
	}
	playTime := time.Now().Sub(starttime)

	if eventCount != len(script) {
		t.Errorf("not all actions were played, want %d, got %d",
			len(script), eventCount)
	}
	want := script[len(script)-1].Time
	if !defaultTimeTolerance.roughlyEqual(playTime, want) {
		t.Errorf("script was not played back at correct speed")
	}
}

func TestPauseResume(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	pauseTime := time.Millisecond * 100

	go func() {
		<-time.After(time.Millisecond * 75)
		if err := p.Pause(); err != nil {
			t.Error(err)
		}
		<-time.After(pauseTime)
		if err := p.Resume(); err != nil {
			t.Error(err)
		}
	}()

	var eventCount int
	starttime := time.Now()
	for range p.Play() {
		eventCount++
	}
	playTime := time.Now().Sub(starttime)

	if eventCount != len(script) {
		t.Errorf("not all actions were played, want %d, got %d",
			len(script), eventCount)
	}
	want := script[len(script)-1].Time + pauseTime
	if !defaultTimeTolerance.roughlyEqual(playTime, want) {
		t.Errorf("script was not played back at correct speed")
	}
}

func TestStop(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	stopTime := time.Millisecond * 75

	go func() {
		<-time.After(stopTime)
		if err := p.Stop(); err != nil {
			t.Error(err)
		}
	}()

	starttime := time.Now()
	for range p.Play() {
		// pass
	}
	playTime := time.Now().Sub(starttime)

	if !defaultTimeTolerance.roughlyEqual(playTime, stopTime) {
		t.Errorf("script was not stopped at the right time, want: %s, got: %s",
			stopTime, playTime)
	}
}

func TestSkip(t *testing.T) {
	cases := []struct {
		Name string
		At   time.Duration
		To   time.Duration
	}{
		{
			Name: "Forward",
			At:   time.Millisecond * 50,
			To:   time.Millisecond * 150,
		},
		{
			Name: "Back",
			At:   time.Millisecond * 50,
			To:   time.Millisecond * 150,
		},
		{
			Name: "Soon",
			At:   time.Millisecond * 0,
			To:   time.Millisecond * 150,
		},
	}

	for _, c := range cases {
		p := NewTimedActionsPlayer()
		p.Script = script

		go func() {
			<-time.After(c.At)
			if err := p.Skip(c.To); err != nil {
				t.Error(err)
			}
		}()

		starttime := time.Now()
		for range p.Play() {
			// pass
		}
		playTime := time.Now().Sub(starttime)
		want := script[len(script)-1].Time - (c.To - c.At)

		if !defaultTimeTolerance.roughlyEqual(playTime, want) {
			t.Errorf("%s: player did not skip correctly, want: %s, got: %s",
				c.Name, want, playTime)
		}
	}
}

func TestLimits(t *testing.T) {
	type TestCase struct {
		Low, High  int
		Slow, Fast int
	}

	cases := []TestCase{
		{0, 100, 0, 0},
		{10, 90, 0, 0},
		{30, 50, 0, 0},
		{50, 90, 0, 0},
		{60, 80, 0, 0},
		{0, 0, 20, 80},
		{0, 0, 30, 60},
		{0, 0, 50, 90},
		{0, 0, 10, 20},
		{0, 0, 0, 50},
		{0, 100, 20, 80},
		{10, 90, 30, 60},
		{30, 50, 50, 90},
		{50, 90, 10, 20},
		{60, 80, 0, 50},
	}
	for i, c := range cases {
		p := NewTimedActionsPlayer()
		p.Script = script
		if c.Low != 0 && c.High != 0 {
			p.LimitPosition(c.Low, c.High)
		}
		if c.Slow != 0 && c.Fast != 0 {
			p.LimitSpeed(c.Slow, c.Fast)
		}

		lowest, highest := c.Low, c.High
		slowest, fastest := c.Slow, c.Fast

		for a := range p.Play() {
			if a.Position < lowest {
				lowest = a.Position
			} else if a.Position > highest {
				highest = a.Position
			}
			if a.Position < slowest {
				slowest = a.Speed
			} else if a.Position > highest {
				fastest = a.Speed
			}
		}

		if c.Low != 0 && c.High != 0 {
			if lowest < c.Low {
				t.Errorf("case %d: went lower than allowed, %d < %d", i, lowest, c.Low)
			}
			if highest > c.High {
				t.Errorf("case %d: went higher than allowed, %d > %d", i, highest, c.High)
			}
		}

		if c.Slow != 0 && c.Fast != 0 {
			if slowest < c.Slow {
				t.Errorf("case %d: went slower than allowed, %d < %d", i, slowest, c.Slow)
			}
			if fastest > c.Fast {
				t.Errorf("case %d: went higher than allowed, %d > %d", i, fastest, c.Fast)
			}
		}

	}
}

func TestDump(t *testing.T) {
	p := NewTimedActionsPlayer()
	p.Script = script

	actions, err := p.Dump()
	if err != nil {
		t.Error(err)
	}
	if len(script) != len(actions) {
		t.Errorf("dump did not return script actions: want %d, got %d",
			len(script), len(actions))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"github.com/growse/pcap"
	"github.com/miekg/dns"
)

const (
	// IPTcp TCP type code
	IPTcp = 6

	// IPUdp UDP type code
	IPUdp = 17
)

// Answer holds dns answer data
type Answer struct {
	Id         int64
	QuestionId int64
	ClientId   int64
	Class      string    `json:"class"`
	Name       string    `json:"name"`
	Record     string    `json:"record"`
	Data       string    `json:"data"`
	Ttl        string    `json:"ttl"`
	CreatedAt  time.Time `json:"created_at"`
}

// Message is used to pass and process data for various output options
type Question struct {
	Id        int64
	SeenCount int64
	ClientId  int64
	Answers   []Answer `json:"answers" sql:"-"`
	Question  string   `json:"question"`
	Length    int      `json:"length" sql:"-"`
	DstIp     string   `json:"dstip"`
	Protocol  string   `json:"protocol" sql:"-"`
	SrcIp     string   `json:"srcip"`
	CreatedAt int64    `json:"created_at"`
	UpdatedAt int64    `json:"updated_at"`
	Packet    string   `json:"packet" sql:"-"`
	Type      uint16   `json:"question_type"`
	Class     uint16   `json:"question_class"`
}

type Client struct {
	Id            int64
	LastSeen      int64
	QuestionCount int64  `json:"question_count"`
	Hostname      string `json:"hostname"`
	Interface     string `json:"interface"`
	MacAddr       string `json:"mac_addr"`
	Ip            string `json:"ip"`
}

// DNS process and parse DNS packets
func DNS(pkt *pcap.Packet, options *Options) (*Question, error) {
	message := &Question{}

	pkt.Decode()
	msg := new(dns.Msg)
	err := msg.Unpack(pkt.Payload)

	if err != nil || len(msg.Answer) <= 0 {
		return message, fmt.Errorf("Error")
	}

	if len(pkt.Headers) <= 0 {
		return message, fmt.Errorf("Error: Missing header information.")
	}

	message.Length = msg.Len()

	packet, _ := msg.Pack()

	message.Packet = hex.EncodeToString(packet)

	ip4hdr, ip4ok := pkt.Headers[0].(*pcap.Iphdr)

	if ip4ok {

		switch ip4hdr.Protocol {
		case IPTcp:
			message.Protocol = "TCP"
		case IPUdp:
			message.Protocol = "UDP"
		default:
			message.Protocol = "N/A"
		}

		message.SrcIp = ip4hdr.SrcAddr()
		message.DstIp = ip4hdr.DestAddr()

	} else {
		ip6hdr, _ := pkt.Headers[0].(*pcap.Ip6hdr)

		message.SrcIp = ip6hdr.SrcAddr()
		message.DstIp = ip6hdr.DestAddr()
		fmt.Println(ip6hdr)
	}

	for i := range msg.Question {
		message.Question = msg.Question[i].Name
		message.Type = msg.Question[i].Qtype
		message.Class = msg.Question[i].Qclass
	}

	for i := range msg.Answer {
		split := strings.Split(msg.Answer[i].String(), "\t")
		answer := Answer{
			Name:      split[0],
			Ttl:       split[1],
			Class:     split[2],
			Record:    split[3],
			Data:      split[4],
			ClientId:  options.Client.Id,
			CreatedAt: time.Now(),
		}

		message.Answers = append(message.Answers, answer)
	}

	return message, nil
}
<commit_msg>remove old print<commit_after>package main

import (
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"github.com/growse/pcap"
	"github.com/miekg/dns"
)

const (
	// IPTcp TCP type code
	IPTcp = 6

	// IPUdp UDP type code
	IPUdp = 17
)

// Answer holds dns answer data
type Answer struct {
	Id         int64
	QuestionId int64
	ClientId   int64
	Class      string    `json:"class"`
	Name       string    `json:"name"`
	Record     string    `json:"record"`
	Data       string    `json:"data"`
	Ttl        string    `json:"ttl"`
	CreatedAt  time.Time `json:"created_at"`
}

// Message is used to pass and process data for various output options
type Question struct {
	Id        int64
	SeenCount int64
	ClientId  int64
	Answers   []Answer `json:"answers" sql:"-"`
	Question  string   `json:"question"`
	Length    int      `json:"length" sql:"-"`
	DstIp     string   `json:"dstip"`
	Protocol  string   `json:"protocol" sql:"-"`
	SrcIp     string   `json:"srcip"`
	CreatedAt int64    `json:"created_at"`
	UpdatedAt int64    `json:"updated_at"`
	Packet    string   `json:"packet" sql:"-"`
	Type      uint16   `json:"question_type"`
	Class     uint16   `json:"question_class"`
}

type Client struct {
	Id            int64
	LastSeen      int64
	QuestionCount int64  `json:"question_count"`
	Hostname      string `json:"hostname"`
	Interface     string `json:"interface"`
	MacAddr       string `json:"mac_addr"`
	Ip            string `json:"ip"`
}

// DNS process and parse DNS packets
func DNS(pkt *pcap.Packet, options *Options) (*Question, error) {
	message := &Question{}

	pkt.Decode()
	msg := new(dns.Msg)
	err := msg.Unpack(pkt.Payload)

	if err != nil || len(msg.Answer) <= 0 {
		return message, fmt.Errorf("Error")
	}

	if len(pkt.Headers) <= 0 {
		return message, fmt.Errorf("Error: Missing header information.")
	}

	message.Length = msg.Len()

	packet, _ := msg.Pack()

	message.Packet = hex.EncodeToString(packet)

	ip4hdr, ip4ok := pkt.Headers[0].(*pcap.Iphdr)

	if ip4ok {

		switch ip4hdr.Protocol {
		case IPTcp:
			message.Protocol = "TCP"
		case IPUdp:
			message.Protocol = "UDP"
		default:
			message.Protocol = "N/A"
		}

		message.SrcIp = ip4hdr.SrcAddr()
		message.DstIp = ip4hdr.DestAddr()

	} else {
		ip6hdr, _ := pkt.Headers[0].(*pcap.Ip6hdr)

		message.SrcIp = ip6hdr.SrcAddr()
		message.DstIp = ip6hdr.DestAddr()
	}

	for i := range msg.Question {
		message.Question = msg.Question[i].Name
		message.Type = msg.Question[i].Qtype
		message.Class = msg.Question[i].Qclass
	}

	for i := range msg.Answer {
		split := strings.Split(msg.Answer[i].String(), "\t")
		answer := Answer{
			Name:      split[0],
			Ttl:       split[1],
			Class:     split[2],
			Record:    split[3],
			Data:      split[4],
			ClientId:  options.Client.Id,
			CreatedAt: time.Now(),
		}

		message.Answers = append(message.Answers, answer)
	}

	return message, nil
}
<|endoftext|>