<commit_before>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 main

import (
	// "encoding/gob"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"

	. "./equtils"

	"github.com/mitchellh/cli"
)

type runConfig struct {
	runScript string
}

func parseRunConfig(jsonPath string) (*runConfig, error) {
	jsonBuf, rerr := WholeRead(jsonPath)
	if rerr != nil {
		return nil, rerr
	}

	var root map[string]interface{}
	err := json.Unmarshal(jsonBuf, &root)
	if err != nil {
		return nil, err
	}

	runScript := root["run"].(string)

	return &runConfig{
		runScript: runScript,
	}, nil
}

func run(args []string) {
	if len(args) != 1 {
		fmt.Printf("specify <storage dir path>\n")
		os.Exit(1)
	}

	storage := args[0]
	confPath := storage + "/" + storageConfigPath

	conf, err := parseRunConfig(confPath)
	if err != nil {
		fmt.Printf("failed to parse config file %s: %s\n", confPath, err)
		os.Exit(1)
	}

	info := readSearchModeDir(storage)
	nextId := info.NrCollectedTraces

	nextDir := storage + "/" + fmt.Sprintf("%08x", nextId)
	err = os.Mkdir(nextDir, 0777)
	if err != nil {
		fmt.Printf("failed to create directory %s: %s\n", nextDir, err)
		os.Exit(1)
	}

	end := make(chan interface{})
	go searchModeNoInitiation(info, nextDir, "random", end)

	materialsDir := storage + "/" + storageMaterialsPath
	runScriptPath := materialsDir + "/" + conf.runScript
	runCmd := exec.Command("sh", "-c", runScriptPath)

	runCmd.Stdout = os.Stdout
	runCmd.Stderr = os.Stderr

	runCmd.Env = append(runCmd.Env, "WORKING_DIR="+nextDir)
	runCmd.Env = append(runCmd.Env, "MATERIALS_DIR="+materialsDir)

	rerr := runCmd.Run()
	if rerr != nil {
		fmt.Printf("failed to execute run script %s: %s\n", runScriptPath, rerr)
		os.Exit(1)
	}

	end <- true
	<-end

	info.NrCollectedTraces++
	updateSearchModeInfo(storage, info)
}

type runCmd struct {
}

func (cmd runCmd) Help() string {
	return "run help (todo)"
}

func (cmd runCmd) Run(args []string) int {
	run(args)
	return 0
}

func (cmd runCmd) Synopsis() string {
	return "run subcommand"
}

func runCommandFactory() (cli.Command, error) {
	return runCmd{}, nil
}
<commit_msg>orchestrator: configurable search policy<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 main

import (
	// "encoding/gob"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"

	. "./equtils"

	"github.com/mitchellh/cli"
)

type runConfig struct {
	runScript    string
	searchPolicy string
}

func parseRunConfig(jsonPath string) (*runConfig, error) {
	jsonBuf, rerr := WholeRead(jsonPath)
	if rerr != nil {
		return nil, rerr
	}

	var root map[string]interface{}
	err := json.Unmarshal(jsonBuf, &root)
	if err != nil {
		return nil, err
	}

	runScript := ""
	if _, ok := root["run"]; ok {
		runScript = root["run"].(string)
	} else {
		fmt.Printf("required field \"run\" is missing\n")
		os.Exit(1)	// TODO: construct suitable error
		return nil, nil
	}

	searchPolicy := "dumb"
	if _, ok := root["searchPolicy"]; ok {
		searchPolicy = root["searchPolicy"].(string)
	}

	return &runConfig{
		runScript: runScript,
		searchPolicy: searchPolicy,
	}, nil
}

func run(args []string) {
	if len(args) != 1 {
		fmt.Printf("specify <storage dir path>\n")
		os.Exit(1)
	}

	storage := args[0]
	confPath := storage + "/" + storageConfigPath

	conf, err := parseRunConfig(confPath)
	if err != nil {
		fmt.Printf("failed to parse config file %s: %s\n", confPath, err)
		os.Exit(1)
	}

	info := readSearchModeDir(storage)
	nextId := info.NrCollectedTraces

	nextDir := storage + "/" + fmt.Sprintf("%08x", nextId)
	err = os.Mkdir(nextDir, 0777)
	if err != nil {
		fmt.Printf("failed to create directory %s: %s\n", nextDir, err)
		os.Exit(1)
	}

	end := make(chan interface{})
	go searchModeNoInitiation(info, nextDir, conf.searchPolicy, end)

	materialsDir := storage + "/" + storageMaterialsPath
	runScriptPath := materialsDir + "/" + conf.runScript
	runCmd := exec.Command("sh", "-c", runScriptPath)

	runCmd.Stdout = os.Stdout
	runCmd.Stderr = os.Stderr

	runCmd.Env = append(runCmd.Env, "WORKING_DIR="+nextDir)
	runCmd.Env = append(runCmd.Env, "MATERIALS_DIR="+materialsDir)

	rerr := runCmd.Run()
	if rerr != nil {
		fmt.Printf("failed to execute run script %s: %s\n", runScriptPath, rerr)
		os.Exit(1)
	}

	end <- true
	<-end

	info.NrCollectedTraces++
	updateSearchModeInfo(storage, info)
}

type runCmd struct {
}

func (cmd runCmd) Help() string {
	return "run help (todo)"
}

func (cmd runCmd) Run(args []string) int {
	run(args)
	return 0
}

func (cmd runCmd) Synopsis() string {
	return "run subcommand"
}

func runCommandFactory() (cli.Command, error) {
	return runCmd{}, nil
}
<|endoftext|>


<commit_before>package end2end_carbon

import (
	"log"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/grafana/metrictank/stacktest/docker"
	"github.com/grafana/metrictank/stacktest/fakemetrics"
	"github.com/grafana/metrictank/stacktest/grafana"
	"github.com/grafana/metrictank/stacktest/graphite"
	"github.com/grafana/metrictank/stacktest/track"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

var tracker *track.Tracker
var fm *fakemetrics.FakeMetrics

const metricsPerSecond = 1000

func TestMain(m *testing.M) {
	log.Println("launching docker-dev stack...")
	cmd := exec.Command("docker-compose", "up", "--force-recreate", "-V")
	cmd.Dir = docker.Path("docker/docker-dev")
	var err error

	tracker, err = track.NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err)
	}

	retcode := m.Run()
	fm.Close()

	log.Println("stopping docker-compose stack...")
	cmd.Process.Signal(syscall.SIGINT)
	// note: even when we don't care about the output, it's best to consume it before calling cmd.Wait()
	// even though the cmd.Wait docs say it will wait for stdout/stderr copying to complete
	// however the docs for cmd.StdoutPipe say "it is incorrect to call Wait before all reads from the pipe have completed"
	tracker.Wait()
	err = cmd.Wait()

	// 130 means ctrl-C (interrupt) which is what we want
	if err != nil && err.Error() != "exit status 130" {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	} else {
		log.Println("docker-compose stack is shut down")
	}

	os.Exit(retcode)
}

func TestStartup(t *testing.T) {
	matchers := []track.Matcher{
		{Str: "metrictank.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank.*carbon-in: listening on.*2003"},
	}
	select {
	case <-tracker.Match(matchers):
		log.Println("stack now running.")
		log.Println("Go to http://localhost:3000 (and login as admin:admin) to see what's going on")
	case <-time.After(time.Second * 70):
		grafana.PostAnnotation("TestStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestBaseIngestWorkload(t *testing.T) {
	grafana.PostAnnotation("TestBaseIngestWorkload:begin")

	fm = fakemetrics.NewCarbon(metricsPerSecond)

	suc6, resp := graphite.RetryGraphite8080("perSecond(metrictank.stats.docker-env.*.input.carbon.metricdata.received.counter32)", "-8s", 18, func(resp graphite.Response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-env.default.input.carbon.metricdata.received.counter32)",
		}
		a := graphite.ValidateTargets(exp)(resp)
		b := graphite.ValidatorLenNulls(1, 8)(resp)
		c := graphite.ValidatorAvgWindowed(8, graphite.Ge(metricsPerSecond))(resp)
		log.Printf("condition target names %t - condition len & nulls %t - condition avg value %t", a, b, c)
		return a && b && c
	})
	if !suc6 {
		grafana.PostAnnotation("TestBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where the MT instance processes at least %d points per second. last response was: %s", metricsPerSecond, spew.Sdump(resp))
	}
}
<commit_msg>show version & fix support for older docker-compose<commit_after>package end2end_carbon

import (
	"log"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/grafana/metrictank/stacktest/docker"
	"github.com/grafana/metrictank/stacktest/fakemetrics"
	"github.com/grafana/metrictank/stacktest/grafana"
	"github.com/grafana/metrictank/stacktest/graphite"
	"github.com/grafana/metrictank/stacktest/track"
)

// TODO: cleanup when ctrl-C go test (teardown all containers)

var tracker *track.Tracker
var fm *fakemetrics.FakeMetrics

const metricsPerSecond = 1000

func TestMain(m *testing.M) {
	log.Println("launching docker-dev stack...")
	version := exec.Command("docker-compose", "version")
	output, err := version.CombinedOutput()
	if err != nil {
		log.Fatal(err.Error())
	}
	log.Println(string(output))

	// TODO: should probably use -V flag here.
	// introduced here https://github.com/docker/compose/releases/tag/1.19.0
	// but circleCI machine image still stuck with 1.14.0
	cmd := exec.Command("docker-compose", "up", "--force-recreate")
	cmd.Dir = docker.Path("docker/docker-dev")

	tracker, err = track.NewTracker(cmd, false, false, "launch-stdout", "launch-stderr")
	if err != nil {
		log.Fatal(err)
	}

	err = cmd.Start()
	if err != nil {
		log.Fatal(err)
	}

	retcode := m.Run()
	fm.Close()

	log.Println("stopping docker-compose stack...")
	cmd.Process.Signal(syscall.SIGINT)
	// note: even when we don't care about the output, it's best to consume it before calling cmd.Wait()
	// even though the cmd.Wait docs say it will wait for stdout/stderr copying to complete
	// however the docs for cmd.StdoutPipe say "it is incorrect to call Wait before all reads from the pipe have completed"
	tracker.Wait()
	err = cmd.Wait()

	// 130 means ctrl-C (interrupt) which is what we want
	if err != nil && err.Error() != "exit status 130" {
		log.Printf("ERROR: could not cleanly shutdown running docker-compose command: %s", err)
		retcode = 1
	} else {
		log.Println("docker-compose stack is shut down")
	}

	os.Exit(retcode)
}

func TestStartup(t *testing.T) {
	matchers := []track.Matcher{
		{Str: "metrictank.*metricIndex initialized.*starting data consumption$"},
		{Str: "metrictank.*carbon-in: listening on.*2003"},
	}
	select {
	case <-tracker.Match(matchers):
		log.Println("stack now running.")
		log.Println("Go to http://localhost:3000 (and login as admin:admin) to see what's going on")
	case <-time.After(time.Second * 70):
		grafana.PostAnnotation("TestStartup:FAIL")
		t.Fatal("timed out while waiting for all metrictank instances to come up")
	}
}

func TestBaseIngestWorkload(t *testing.T) {
	grafana.PostAnnotation("TestBaseIngestWorkload:begin")

	fm = fakemetrics.NewCarbon(metricsPerSecond)

	suc6, resp := graphite.RetryGraphite8080("perSecond(metrictank.stats.docker-env.*.input.carbon.metricdata.received.counter32)", "-8s", 18, func(resp graphite.Response) bool {
		exp := []string{
			"perSecond(metrictank.stats.docker-env.default.input.carbon.metricdata.received.counter32)",
		}
		a := graphite.ValidateTargets(exp)(resp)
		b := graphite.ValidatorLenNulls(1, 8)(resp)
		c := graphite.ValidatorAvgWindowed(8, graphite.Ge(metricsPerSecond))(resp)
		log.Printf("condition target names %t - condition len & nulls %t - condition avg value %t", a, b, c)
		return a && b && c
	})
	if !suc6 {
		grafana.PostAnnotation("TestBaseIngestWorkload:FAIL")
		t.Fatalf("cluster did not reach a state where the MT instance processes at least %d points per second. last response was: %s", metricsPerSecond, spew.Sdump(resp))
	}
}
<|endoftext|>


<commit_before>package packer

import "fmt"

// The version of packer.
const Version = "0.1.0.dev"

type versionCommand byte

func (versionCommand) Help() string {
	return `usage: packer version

Outputs the version of Packer that is running. There are no additional
command-line flags for this command.`
}

func (versionCommand) Run(env Environment, args []string) int {
	env.Ui().Say(fmt.Sprintf("Packer v%v\n", Version))
	return 0
}

func (versionCommand) Synopsis() string {
	return "print Packer version"
}
<commit_msg>packer: Remove extra newline on version<commit_after>package packer

import "fmt"

// The version of packer.
const Version = "0.1.0.dev"

type versionCommand byte

func (versionCommand) Help() string {
	return `usage: packer version

Outputs the version of Packer that is running. There are no additional
command-line flags for this command.`
}

func (versionCommand) Run(env Environment, args []string) int {
	env.Ui().Say(fmt.Sprintf("Packer v%v", Version))
	return 0
}

func (versionCommand) Synopsis() string {
	return "print Packer version"
}
<|endoftext|>


<commit_before>package firebase

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"github.com/knq/jwt"
	"github.com/knq/oauth2util"
)

const (
	// DefaultTokenExpiration is the default expiration for generated OAuth2 tokens.
	DefaultTokenExpiration = 1 * time.Hour
)

// Option is an option to modify a Firebase ref.
type Option func(r *Ref) error

// URL is an option to set Firebase base URL.
func URL(urlstr string) Option {
	return func(r *Ref) error {
		u, err := url.Parse(urlstr)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not parse url: %v", err),
			}
		}

		r.url = u

		return nil
	}
}

// Transport is an option to set the HTTP transport.
func Transport(roundTripper http.RoundTripper) Option {
	return func(r *Ref) error {
		r.transport = roundTripper
		return nil
	}
}

// WatchBufferLen is an option that sets the watch buffer size.
func WatchBufferLen(len int) Option {
	return func(r *Ref) error {
		r.watchBufLen = len
		return nil
	}
}

// GoogleServiceAccountCredentialsJSON loads Firebase credentials from the JSON
// encoded buf.
func GoogleServiceAccountCredentialsJSON(buf []byte) Option {
	return func(r *Ref) error {
		var err error

		var v struct {
			ProjectID   string `json:"project_id"`
			ClientEmail string `json:"client_email"`
			PrivateKey  string `json:"private_key"`
			TokenURI    string `json:"token_uri"`
		}

		// decode settings into v
		err = json.Unmarshal(buf, &v)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not unmarshal service account credentials: %v", err),
			}
		}

		// simple check
		if v.ProjectID == "" || v.ClientEmail == "" || v.PrivateKey == "" {
			return &Error{
				Err: "google service account credentials missing project_id, client_email or private_key",
			}
		}

		// set URL
		err = URL("https://" + v.ProjectID + ".firebaseio.com/")(r)
		if err != nil {
			return err
		}

		// create token signer
		signer, err := jwt.RS256.New(jwt.PEM{[]byte(v.PrivateKey)})
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not create jwt signer for auth token source: %v", err),
			}
		}

		// create auth token source
		r.auth, err = oauth2util.JWTBearerGrantTokenSource(
			signer, v.TokenURI, context.Background(),
			oauth2util.ExpiresIn(DefaultTokenExpiration),
			oauth2util.IssuedAt(true),
		)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not create auth token source: %v", err),
			}
		}

		// add the claims for firebase
		r.auth.AddClaim("iss", v.ClientEmail)
		r.auth.AddClaim("sub", v.ClientEmail)
		r.auth.AddClaim("aud", v.TokenURI)
		r.auth.AddClaim("scope", strings.Join([]string{
			"https://www.googleapis.com/auth/userinfo.email",
			"https://www.googleapis.com/auth/firebase.database",
		}, " "))

		// set source
		r.source = oauth2.ReuseTokenSource(nil, r.auth)

		return nil
	}
}

// GoogleServiceAccountCredentialsFile loads Firebase credentials from the
// specified path on disk and configures Firebase accordingly.
//
// Account credentials can be downloaded from the Google Cloud console.
func GoogleServiceAccountCredentialsFile(path string) Option {
	return func(r *Ref) error {
		buf, err := ioutil.ReadFile(path)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not read google service account credentials file: %v", err),
			}
		}

		return GoogleServiceAccountCredentialsJSON(buf)(r)
	}
}

// DefaultQueryOptions is an option that sets the default query options on the
// ref.
func DefaultQueryOptions(opts ...QueryOption) Option {
	return func(r *Ref) error {
		r.SetQueryOptions(opts...)
		return nil
	}
}

// UserID is an option that sets the auth user id ("uid") via the
// auth_variable_override on the ref.
func UserID(uid string) Option {
	return func(r *Ref) error {
		return DefaultQueryOptions(AuthOverride(map[string]interface{}{
			"uid": uid,
		}))(r)
	}
}

// WithClaims is an option that adds additional claims to the auth token source.
func WithClaims(claims map[string]interface{}) Option {
	return func(r *Ref) error {
		return r.AddTokenSourceClaim("claims", claims)
	}
}

// httpLogger handles logging http requests and responses.
type httpLogger struct {
	transport                 http.RoundTripper
	requestLogf, responseLogf Logf
}

// RoundTrip satisifies the http.RoundTripper interface.
func (hl *httpLogger) RoundTrip(req *http.Request) (*http.Response, error) {
	trans := hl.transport
	if trans == nil {
		trans = http.DefaultTransport
	}

	reqBody, _ := httputil.DumpRequestOut(req, true)
	res, err := trans.RoundTrip(req)
	resBody, _ := httputil.DumpResponse(res, true)

	hl.requestLogf("%s", reqBody)
	hl.responseLogf("%s", resBody)

	return res, err
}

// Logf is a logging func.
type Logf func(string, ...interface{})

// Log is an option that writes all HTTP request and response data to the
// respective logger.
func Log(requestLogf, responseLogf Logf) Option {
	return func(r *Ref) error {
		return Transport(&httpLogger{
			transport:    r.transport,
			requestLogf:  requestLogf,
			responseLogf: responseLogf,
		})(r)
	}
}

// QueryOption is an option used to modify the underlying http.Request for
// Firebase.
type QueryOption func(url.Values) error

// Shallow is a query option that toggles Firebase to return a shallow result.
func Shallow(v url.Values) error {
	v.Add("shallow", "true")
	return nil
}

// PrintPretty is a query option that toggles pretty formatting for Firebase
// results.
func PrintPretty(v url.Values) error {
	v.Add("print", "pretty")
	return nil
}

// jsonQuery returns a QueryOption for a field and json encodes the val.
func jsonQuery(field string, val interface{}) QueryOption {
	// json encode
	buf, err := json.Marshal(val)
	if err != nil {
		err = &Error{
			Err: fmt.Sprintf("could not marshal query option: %v", err),
		}
	}

	return func(v url.Values) error {
		if err != nil {
			return err
		}

		v.Add(field, string(buf))
		return nil
	}
}

// uintQuery returns a QueryOption for a field that converts n into a string.
func uintQuery(field string, n uint) QueryOption {
	val := strconv.FormatUint(uint64(n), 10)
	return func(v url.Values) error {
		v.Add(field, val)
		return nil
	}
}

// OrderBy is a query option that sets Firebase's returned result order.
func OrderBy(field string) QueryOption {
	return jsonQuery("orderBy", field)
}

// EqualTo is a query option that sets the order by filter to equalTo val.
func EqualTo(val interface{}) QueryOption {
	return jsonQuery("equalTo", val)
}

// StartAt is a query option that sets the order by filter to startAt val.
func StartAt(val interface{}) QueryOption {
	return jsonQuery("startAt", val)
}

// EndAt is a query option that sets the order by filter to endAt val.
func EndAt(val interface{}) QueryOption {
	return jsonQuery("endAt", val)
}

// AuthOverride is a query option that sets the auth_variable_override.
func AuthOverride(val interface{}) QueryOption {
	return jsonQuery("auth_variable_override", val)
}

// LimitToFirst is a query option that limit's Firebase's returned results to
// the first n items.
func LimitToFirst(n uint) QueryOption {
	return uintQuery("limitToFirst", n)
}

// LimitToLast is a query option that limit's Firebase's returned results to
// the last n items.
func LimitToLast(n uint) QueryOption {
	return uintQuery("limitToLast", n)
}
<commit_msg>Quick code cleanup<commit_after>package firebase

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"github.com/knq/jwt"
	"github.com/knq/oauth2util"
)

const (
	// DefaultTokenExpiration is the default expiration for generated OAuth2 tokens.
	DefaultTokenExpiration = 1 * time.Hour
)

// Option is an option to modify a Firebase ref.
type Option func(r *Ref) error

// URL is an option to set Firebase base URL.
func URL(urlstr string) Option {
	return func(r *Ref) error {
		u, err := url.Parse(urlstr)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not parse url: %v", err),
			}
		}

		r.url = u

		return nil
	}
}

// Transport is an option to set the HTTP transport.
func Transport(roundTripper http.RoundTripper) Option {
	return func(r *Ref) error {
		r.transport = roundTripper
		return nil
	}
}

// WatchBufferLen is an option that sets the watch buffer size.
func WatchBufferLen(len int) Option {
	return func(r *Ref) error {
		r.watchBufLen = len
		return nil
	}
}

// GoogleServiceAccountCredentialsJSON loads Firebase credentials from the JSON
// encoded buf.
func GoogleServiceAccountCredentialsJSON(buf []byte) Option {
	return func(r *Ref) error {
		var err error

		var v struct {
			ProjectID   string `json:"project_id"`
			ClientEmail string `json:"client_email"`
			PrivateKey  string `json:"private_key"`
			TokenURI    string `json:"token_uri"`
		}

		// decode settings into v
		err = json.Unmarshal(buf, &v)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not unmarshal service account credentials: %v", err),
			}
		}

		// simple check
		if v.ProjectID == "" || v.ClientEmail == "" || v.PrivateKey == "" {
			return &Error{
				Err: "google service account credentials missing project_id, client_email or private_key",
			}
		}

		// set URL
		err = URL("https://" + v.ProjectID + ".firebaseio.com/")(r)
		if err != nil {
			return err
		}

		// create token signer
		signer, err := jwt.RS256.New(jwt.PEM{[]byte(v.PrivateKey)})
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not create jwt signer for auth token source: %v", err),
			}
		}

		// create auth token source
		r.auth, err = oauth2util.JWTBearerGrantTokenSource(
			signer, v.TokenURI, context.Background(),
			oauth2util.ExpiresIn(DefaultTokenExpiration),
			oauth2util.IssuedAt(true),
		)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not create auth token source: %v", err),
			}
		}

		// add the claims for firebase
		r.auth.AddClaim("iss", v.ClientEmail)
		r.auth.AddClaim("sub", v.ClientEmail)
		r.auth.AddClaim("aud", v.TokenURI)
		r.auth.AddClaim("scope", strings.Join([]string{
			"https://www.googleapis.com/auth/userinfo.email",
			"https://www.googleapis.com/auth/firebase.database",
		}, " "))

		// set source
		r.source = oauth2.ReuseTokenSource(nil, r.auth)

		return nil
	}
}

// GoogleServiceAccountCredentialsFile loads Firebase credentials from the
// specified path on disk and configures Firebase accordingly.
//
// Account credentials can be downloaded from the Google Cloud console.
func GoogleServiceAccountCredentialsFile(path string) Option {
	return func(r *Ref) error {
		buf, err := ioutil.ReadFile(path)
		if err != nil {
			return &Error{
				Err: fmt.Sprintf("could not read google service account credentials file: %v", err),
			}
		}

		return GoogleServiceAccountCredentialsJSON(buf)(r)
	}
}

// DefaultQueryOptions is an option that sets the default query options on the
// ref.
func DefaultQueryOptions(opts ...QueryOption) Option {
	return func(r *Ref) error {
		r.SetQueryOptions(opts...)
		return nil
	}
}

// UserID is an option that sets the auth user id ("uid") via the
// auth_variable_override on the ref.
func UserID(uid string) Option {
	return func(r *Ref) error {
		return DefaultQueryOptions(
			AuthOverride(map[string]interface{}{
				"uid": uid,
			}),
		)(r)
	}
}

// WithClaims is an option that adds additional claims to the auth token source.
func WithClaims(claims map[string]interface{}) Option {
	return func(r *Ref) error {
		return r.AddTokenSourceClaim("claims", claims)
	}
}

// httpLogger handles logging http requests and responses.
type httpLogger struct {
	transport                 http.RoundTripper
	requestLogf, responseLogf Logf
}

// RoundTrip satisifies the http.RoundTripper interface.
func (hl *httpLogger) RoundTrip(req *http.Request) (*http.Response, error) {
	trans := hl.transport
	if trans == nil {
		trans = http.DefaultTransport
	}

	reqBody, _ := httputil.DumpRequestOut(req, true)
	res, err := trans.RoundTrip(req)
	resBody, _ := httputil.DumpResponse(res, true)

	hl.requestLogf("%s", reqBody)
	hl.responseLogf("%s", resBody)

	return res, err
}

// Logf is a logging func.
type Logf func(string, ...interface{})

// Log is an option that writes all HTTP request and response data to the
// respective logger.
func Log(requestLogf, responseLogf Logf) Option {
	return func(r *Ref) error {
		return Transport(&httpLogger{
			transport:    r.transport,
			requestLogf:  requestLogf,
			responseLogf: responseLogf,
		})(r)
	}
}

// QueryOption is an option used to modify the underlying http.Request for
// Firebase.
type QueryOption func(url.Values) error

// Shallow is a query option that toggles Firebase to return a shallow result.
func Shallow(v url.Values) error {
	v.Add("shallow", "true")
	return nil
}

// PrintPretty is a query option that toggles pretty formatting for Firebase
// results.
func PrintPretty(v url.Values) error {
	v.Add("print", "pretty")
	return nil
}

// jsonQuery returns a QueryOption for a field and json encodes the val.
func jsonQuery(field string, val interface{}) QueryOption {
	// json encode
	buf, err := json.Marshal(val)
	if err != nil {
		err = &Error{
			Err: fmt.Sprintf("could not marshal query option: %v", err),
		}
	}

	return func(v url.Values) error {
		if err != nil {
			return err
		}

		v.Add(field, string(buf))
		return nil
	}
}

// uintQuery returns a QueryOption for a field that converts n into a string.
func uintQuery(field string, n uint) QueryOption {
	val := strconv.FormatUint(uint64(n), 10)
	return func(v url.Values) error {
		v.Add(field, val)
		return nil
	}
}

// OrderBy is a query option that sets Firebase's returned result order.
func OrderBy(field string) QueryOption {
	return jsonQuery("orderBy", field)
}

// EqualTo is a query option that sets the order by filter to equalTo val.
func EqualTo(val interface{}) QueryOption {
	return jsonQuery("equalTo", val)
}

// StartAt is a query option that sets the order by filter to startAt val.
func StartAt(val interface{}) QueryOption {
	return jsonQuery("startAt", val)
}

// EndAt is a query option that sets the order by filter to endAt val.
func EndAt(val interface{}) QueryOption {
	return jsonQuery("endAt", val)
}

// AuthOverride is a query option that sets the auth_variable_override.
func AuthOverride(val interface{}) QueryOption {
	return jsonQuery("auth_variable_override", val)
}

// LimitToFirst is a query option that limit's Firebase's returned results to
// the first n items.
func LimitToFirst(n uint) QueryOption {
	return uintQuery("limitToFirst", n)
}

// LimitToLast is a query option that limit's Firebase's returned results to
// the last n items.
func LimitToLast(n uint) QueryOption {
	return uintQuery("limitToLast", n)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"sort"
)

const packageTemplateString = `<!DOCTYPE html>

		<meta charset="utf-8">
		<title>{{.Repo.PackageName}}.{{.Repo.MajorVersion}}{{.Repo.SubPath}} - {{.Repo.GopkgPath}}</title>
		<link href='//fonts.googleapis.com/css?family=Ubuntu+Mono|Ubuntu' rel='stylesheet' >
		<link href="//netdna.bootstrapcdn.com/font-awesome/4.0.3/css/font-awesome.css" rel="stylesheet" >
		<link href="//netdna.bootstrapcdn.com/bootstrap/3.1.1/css/bootstrap.min.css" rel="stylesheet" >

{{.Repo.GopkgPath}}

						 Source Code
						 API Documentation

Getting started

To get the package, execute:

go get {{.Repo.GopkgPath}}

To import this package, add the following line to your code:

import "{{.Repo.GopkgPath}}"

								{{if .CleanPackageName}}
Refer to it as {{.Repo.PackageName}}.{{end}}

For more details, see the API documentation.

Versions

						{{ if .LatestVersions }}
							{{ range .LatestVersions }}

									v{{.Major}}
									→
									{{.}}

							{{ end }}
						{{ else }}

								v0
								→
								master

						{{ end }}

gopkg.in

`

var packageTemplate *template.Template

func gopkgVersionRoot(repo *Repo, version Version) string {
	return repo.GopkgVersionRoot(version)
}

var packageFuncs = template.FuncMap{
	"gopkgVersionRoot": gopkgVersionRoot,
}

func init() {
	var err error
	packageTemplate, err = template.New("page").Funcs(packageFuncs).Parse(packageTemplateString)
	if err != nil {
		fmt.Fprintf(os.Stderr, "fatal: parsing package template failed: %s\n", err)
		os.Exit(1)
	}
}

type packageData struct {
	Repo             *Repo
	LatestVersions   VersionList // Contains only the latest version for each major
	FullVersion      Version     // Version that the major requested resolves to
	CleanPackageName bool
}

func renderPackagePage(resp http.ResponseWriter, req *http.Request, repo *Repo) {
	data := &packageData{
		Repo: repo,
	}
	latestVersionsMap := make(map[int]Version)
	for _, v := range repo.AllVersions {
		v2, exists := latestVersionsMap[v.Major]
		if !exists || v2.Less(v) {
			latestVersionsMap[v.Major] = v
		}
	}
	data.FullVersion = latestVersionsMap[repo.MajorVersion.Major]
	data.LatestVersions = make(VersionList, 0, len(latestVersionsMap))
	for _, v := range latestVersionsMap {
		data.LatestVersions = append(data.LatestVersions, v)
	}
	sort.Sort(sort.Reverse(data.LatestVersions))

	data.CleanPackageName = true
	for i, c := range repo.PackageName {
		if c < 'a' || c > 'z' {
			if i == 0 || c != '_' && (c < '0' || c > '9') {
				data.CleanPackageName = false
				break
			}
		}
	}

	err := packageTemplate.Execute(resp, data)
	if err != nil {
		log.Printf("error executing tmplPackage: %s\n", err)
	}
}
<commit_msg>Strip out go- and -go from clean package name.<commit_after>package main

import (
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
)

const packageTemplateString = `<!DOCTYPE html>

		<meta charset="utf-8">
		<title>{{.Repo.PackageName}}.{{.Repo.MajorVersion}}{{.Repo.SubPath}} - {{.Repo.GopkgPath}}</title>
		<link href='//fonts.googleapis.com/css?family=Ubuntu+Mono|Ubuntu' rel='stylesheet' >
		<link href="//netdna.bootstrapcdn.com/font-awesome/4.0.3/css/font-awesome.css" rel="stylesheet" >
		<link href="//netdna.bootstrapcdn.com/bootstrap/3.1.1/css/bootstrap.min.css" rel="stylesheet" >

{{.Repo.GopkgPath}}

						 Source Code
						 API Documentation

Getting started

To get the package, execute:

go get {{.Repo.GopkgPath}}

To import this package, add the following line to your code:

import "{{.Repo.GopkgPath}}"

								{{if .CleanPackageName}}
Refer to it as {{.CleanPackageName}}.{{end}}

For more details, see the API documentation.

Versions

						{{ if .LatestVersions }}
							{{ range .LatestVersions }}

									v{{.Major}}
									→
									{{.}}

							{{ end }}
						{{ else }}

								v0
								→
								master

						{{ end }}

gopkg.in

`

var packageTemplate *template.Template

func gopkgVersionRoot(repo *Repo, version Version) string {
	return repo.GopkgVersionRoot(version)
}

var packageFuncs = template.FuncMap{
	"gopkgVersionRoot": gopkgVersionRoot,
}

func init() {
	var err error
	packageTemplate, err = template.New("page").Funcs(packageFuncs).Parse(packageTemplateString)
	if err != nil {
		fmt.Fprintf(os.Stderr, "fatal: parsing package template failed: %s\n", err)
		os.Exit(1)
	}
}

type packageData struct {
	Repo             *Repo
	LatestVersions   VersionList // Contains only the latest version for each major
	FullVersion      Version     // Version that the major requested resolves to
	CleanPackageName string
}

func renderPackagePage(resp http.ResponseWriter, req *http.Request, repo *Repo) {
	data := &packageData{
		Repo: repo,
	}
	latestVersionsMap := make(map[int]Version)
	for _, v := range repo.AllVersions {
		v2, exists := latestVersionsMap[v.Major]
		if !exists || v2.Less(v) {
			latestVersionsMap[v.Major] = v
		}
	}
	data.FullVersion = latestVersionsMap[repo.MajorVersion.Major]
	data.LatestVersions = make(VersionList, 0, len(latestVersionsMap))
	for _, v := range latestVersionsMap {
		data.LatestVersions = append(data.LatestVersions, v)
	}
	sort.Sort(sort.Reverse(data.LatestVersions))

	data.CleanPackageName = repo.PackageName
	if strings.HasPrefix(data.CleanPackageName, "go-") {
		data.CleanPackageName = data.CleanPackageName[3:]
	}
	if strings.HasSuffix(data.CleanPackageName, "-go") {
		data.CleanPackageName = data.CleanPackageName[:len(data.CleanPackageName)-3]
	}
	for i, c := range data.CleanPackageName {
		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' {
			continue
		}
		if i > 0 && (c == '_' || c >= '0' && c <= '9') {
			continue
		}
		data.CleanPackageName = ""
		break
	}

	err := packageTemplate.Execute(resp, data)
	if err != nil {
		log.Printf("error executing tmplPackage: %s\n", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 acceptance

import (
	"path/filepath"
	"testing"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/acceptance"
)

func init() {
	acceptance.DefineFlags()
}

func TestAcceptanceDotNet(t *testing.T) {
	builder, cleanup := acceptance.CreateBuilder(t)
	defer cleanup()

	sdk := filepath.Join("/layers", dotnetRuntime, "sdk")

	testCases := []acceptance.Test{
		{
			Name:              "simple dotnet app",
			App:               "dotnet/simple",
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple dotnet app with runtime version",
			App:               "dotnet/simple",
			Path:              "/version?want=3.1.1",
			Env:               []string{"GOOGLE_RUNTIME_VERSION=3.1.101"},
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple dotnet app with global.json",
			App:               "dotnet/simple_with_global",
			Path:              "/version?want=3.1.0",
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple prebuilt dotnet app",
			App:               "dotnet/simple_prebuilt",
			Env:               []string{"GOOGLE_ENTRYPOINT=./simple"},
			MustUse:           []string{dotnetRuntime},
			MustNotUse:        []string{dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:    "simple dotnet app (Dev Mode)",
			App:     "dotnet/simple",
			Env:     []string{"GOOGLE_DEVMODE=1"},
			MustUse: []string{dotnetRuntime, dotnetPublish},
		},
		{
			Name:       "dotnet selected via GOOGLE_RUNTIME",
			App:        "override",
			Env:        []string{"GOOGLE_RUNTIME=dotnet"},
			MustUse:    []string{dotnetRuntime},
			MustNotUse: []string{nodeRuntime, pythonRuntime, goRuntime},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			acceptance.TestApp(t, builder, tc)
		})
	}
}

func TestFailuresDotNet(t *testing.T) {
	builder, cleanup := acceptance.CreateBuilder(t)
	t.Cleanup(cleanup)

	testCases := []acceptance.FailureTest{
		{
			Name:      "bad runtime version",
			App:       "dotnet/simple",
			Env:       []string{"GOOGLE_RUNTIME_VERSION=BAD_NEWS_BEARS"},
			MustMatch: "Runtime version BAD_NEWS_BEARS does not exist",
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.Name, func(t *testing.T) {
			t.Parallel()

			acceptance.TestBuildFailure(t, builder, tc)
		})
	}
}
<commit_msg>Check MustRebuildOnChange in .NET Dev Mode test.<commit_after>// Copyright 2020 Google LLC
//
// 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 acceptance

import (
	"path/filepath"
	"testing"

	"github.com/GoogleCloudPlatform/buildpacks/pkg/acceptance"
)

func init() {
	acceptance.DefineFlags()
}

func TestAcceptanceDotNet(t *testing.T) {
	builder, cleanup := acceptance.CreateBuilder(t)
	defer cleanup()

	sdk := filepath.Join("/layers", dotnetRuntime, "sdk")

	testCases := []acceptance.Test{
		{
			Name:              "simple dotnet app",
			App:               "dotnet/simple",
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple dotnet app with runtime version",
			App:               "dotnet/simple",
			Path:              "/version?want=3.1.1",
			Env:               []string{"GOOGLE_RUNTIME_VERSION=3.1.101"},
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple dotnet app with global.json",
			App:               "dotnet/simple_with_global",
			Path:              "/version?want=3.1.0",
			MustUse:           []string{dotnetRuntime, dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:              "simple prebuilt dotnet app",
			App:               "dotnet/simple_prebuilt",
			Env:               []string{"GOOGLE_ENTRYPOINT=./simple"},
			MustUse:           []string{dotnetRuntime},
			MustNotUse:        []string{dotnetPublish},
			FilesMustNotExist: []string{sdk},
		},
		{
			Name:                "simple dotnet app (Dev Mode)",
			App:                 "dotnet/simple",
			Env:                 []string{"GOOGLE_DEVMODE=1"},
			MustUse:             []string{dotnetRuntime, dotnetPublish},
			FilesMustExist:      []string{sdk, "/workspace/Startup.cs"},
			MustRebuildOnChange: "/workspace/Startup.cs",
		},
		{
			Name:       "dotnet selected via GOOGLE_RUNTIME",
			App:        "override",
			Env:        []string{"GOOGLE_RUNTIME=dotnet"},
			MustUse:    []string{dotnetRuntime},
			MustNotUse: []string{nodeRuntime, pythonRuntime, goRuntime},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.Name, func(t *testing.T) {
			acceptance.TestApp(t, builder, tc)
		})
	}
}

func TestFailuresDotNet(t *testing.T) {
	builder, cleanup := acceptance.CreateBuilder(t)
	t.Cleanup(cleanup)

	testCases := []acceptance.FailureTest{
		{
			Name:      "bad runtime version",
			App:       "dotnet/simple",
			Env:       []string{"GOOGLE_RUNTIME_VERSION=BAD_NEWS_BEARS"},
			MustMatch: "Runtime version BAD_NEWS_BEARS does not exist",
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.Name, func(t *testing.T) {
			t.Parallel()

			acceptance.TestBuildFailure(t, builder, tc)
		})
	}
}
<|endoftext|>


<commit_before>package rain

import (
	"bytes"
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/protocol"
)

// All current implementations use 2^14 (16 kiB), and close connections which request an amount greater than that.
const blockSize = 16 * 1024

type peerConn struct {
	conn         net.Conn
	disconnected chan struct{} // will be closed when peer disconnects

	unchokeM       sync.Mutex    // protects unchokeC
	unchokeC       chan struct{} // will be closed when and "unchoke" message is received
	onceInterested sync.Once     // for sending "interested" message only once

	amChoking      bool // this client is choking the peer
	amInterested   bool // this client is interested in the peer
	peerChoking    bool // peer is choking this client
	peerInterested bool // peer is interested in this client
	// peerRequests   map[uint64]bool      // What remote peer requested
	// ourRequests    map[uint64]time.Time // What we requested, when we requested it

	log logger.Logger
}

func newPeerConn(conn net.Conn) *peerConn {
	return &peerConn{
		conn:         conn,
		disconnected: make(chan struct{}),
		unchokeC:     make(chan struct{}),
		amChoking:    true,
		peerChoking:  true,
		log:          logger.New("peer " + conn.RemoteAddr().String()),
	}
}

const connReadTimeout = 3 * time.Minute

// run processes incoming messages after handshake.
func (p *peerConn) run(t *transfer) {
	defer close(p.disconnected)
	p.log.Debugln("Communicating peer", p.conn.RemoteAddr())

	bitField := bitfield.New(nil, t.bitField.Len())

	err := p.sendBitField(t.bitField)
	if err != nil {
		p.log.Error(err)
		return
	}

	first := true
	for {
		err = p.conn.SetReadDeadline(time.Now().Add(connReadTimeout))
		if err != nil {
			p.log.Error(err)
			return
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			if err == io.EOF {
				p.log.Warning("Remote peer has closed the connection")
				return
			}
			p.log.Error(err)
			return
		}
		p.log.Debugf("Received message of length: %d", length)

		if length == 0 { // keep-alive message
			p.log.Debug("Received message of type \"keep alive\"")
			continue
		}

		var msgType protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &msgType)
		if err != nil {
			p.log.Error(err)
			return
		}
		length--

		p.log.Debugf("Received message of type %q", msgType)

		switch msgType {
		case protocol.Choke:
			p.peerChoking = true
		case protocol.Unchoke:
			p.unchokeM.Lock()
			p.peerChoking = false
			close(p.unchokeC)
			// p.unchokeC = make(chan struct{})
			p.unchokeM.Unlock()
		case protocol.Interested:
			p.peerInterested = true
		case protocol.NotInterested:
			p.peerInterested = false
		case protocol.Have:
			var i uint32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				p.log.Error(err)
				return
			}
			if i >= uint32(len(t.pieces)) {
				p.log.Error("unexpected piece index")
				return
			}
			piece := t.pieces[i]
			bitField.Set(i)
			p.log.Debug("Peer ", p.conn.RemoteAddr(), " has piece #", i)
			p.log.Debugln("new bitfield:", bitField.Hex())

			t.haveC <- peerHave{p, piece}
		case protocol.Bitfield:
			if !first {
				p.log.Error("bitfield can only be sent after handshake")
				return
			}

			if int64(length) != int64(len(bitField.Bytes())) {
				p.log.Error("invalid bitfield length")
				return
			}

			_, err = p.conn.Read(bitField.Bytes())
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugln("Received bitfield:", bitField.Hex())

			for i := uint32(0); i < bitField.Len(); i++ {
				if bitField.Test(i) {
					t.haveC <- peerHave{p, t.pieces[i]}
				}
			}
		case protocol.Request:
		case protocol.Piece:
			var index uint32
			err = binary.Read(p.conn, binary.BigEndian, &index)
			if err != nil {
				p.log.Error(err)
				return
			}
			if index >= uint32(len(t.pieces)) {
				p.log.Error("unexpected piece index")
				return
			}
			piece := t.pieces[index]
			var begin uint32
			err = binary.Read(p.conn, binary.BigEndian, &begin)
			if err != nil {
				p.log.Error(err)
				return
			}
			if begin%blockSize != 0 {
				p.log.Error("unexpected piece offset")
				return
			}
			blockIndex := begin / blockSize
			if blockIndex >= uint32(len(piece.blocks)) {
				p.log.Error("unexpected piece offset")
				return
			}
			block := &piece.blocks[blockIndex]
			length -= 8
			if length != block.length {
				p.log.Error("unexpected block size")
				return
			}
			data := make([]byte, length)
			_, err = io.ReadFull(p.conn, data)
			if err != nil {
				p.log.Error(err)
				return
			}
			piece.blockC <- peerBlock{p, block, data}
		case protocol.Cancel:
		case protocol.Port:
		default:
			p.log.Debugf("Unknown message type: %d", msgType)
			// Discard remaining bytes.
			io.CopyN(ioutil.Discard, p.conn, int64(length))
		}

		first = false
	}
}

func (p *peerConn) sendBitField(b bitfield.BitField) error {
	var buf bytes.Buffer
	length := int32(1 + len(b.Bytes()))
	buf.Grow(4 + int(length))
	err := binary.Write(&buf, binary.BigEndian, length)
	if err != nil {
		return err
	}
	if err = buf.WriteByte(byte(protocol.Bitfield)); err != nil {
		return err
	}
	if _, err = buf.Write(b.Bytes()); err != nil {
		return err
	}
	p.log.Debugf("Sending message: \"bitfield\" %#v", buf.Bytes())
	_, err = io.Copy(p.conn, &buf)
	return err
}

// beInterested sends "interested" message to peer (once) and
// returns a channel that will be closed when an "unchoke" message is received.
func (p *peerConn) beInterested() (unchokeC chan struct{}, err error) {
	p.log.Debug("beInterested")
	p.unchokeM.Lock()
	defer p.unchokeM.Unlock()

	unchokeC = p.unchokeC

	if !p.peerChoking {
		return
	}

	p.onceInterested.Do(func() { err = p.sendMessage(protocol.Interested) })
	return
}

func (p *peerConn) sendMessage(msgType protocol.MessageType) error {
	var msg = struct {
		Length      uint32
		MessageType protocol.MessageType
	}{1, msgType}
	p.log.Debugf("Sending message: %q", msgType)
	return binary.Write(p.conn, binary.BigEndian, &msg)
}

type peerRequestMessage struct {
	ID                   protocol.MessageType
	Index, Begin, Length uint32
}

func newPeerRequestMessage(index, begin, length uint32) *peerRequestMessage {
	return &peerRequestMessage{protocol.Request, index, begin, length}
}

func (p *peerConn) sendRequest(m *peerRequestMessage) error {
	var msg = struct {
		Length  uint32
		Message peerRequestMessage
	}{13, *m}
	p.log.Debugf("Sending message: %q %#v", "request", msg)
	return binary.Write(p.conn, binary.BigEndian, &msg)
}

func (p *peerConn) downloadPiece(piece *piece) error {
	p.log.Debugf("downloading piece #%d", piece.index)

	unchokeC, err := p.beInterested()
	if err != nil {
		return err
	}

	select {
	case <-unchokeC:
	case <-time.After(time.Minute):
		p.conn.Close()
		return errors.New("Peer did not unchoke")
	}

	for _, b := range piece.blocks {
		if err := p.sendRequest(newPeerRequestMessage(piece.index, b.index*blockSize, b.length)); err != nil {
			return err
		}
	}

	pieceData := make([]byte, piece.length)
	for _ = range piece.blocks {
		select {
		case peerBlock := <-piece.blockC:
			p.log.Infoln("received block of length", len(peerBlock.data))
			copy(pieceData[peerBlock.block.index*blockSize:], peerBlock.data)
			if _, err = peerBlock.block.files.Write(peerBlock.data); err != nil {
				return err
			}
			piece.bitField.Set(peerBlock.block.index)
		case <-time.After(time.Minute):
			return fmt.Errorf("peer did not send piece #%d completely", piece.index)
		}
	}

	// Verify piece hash
	hash := sha1.New()
	hash.Write(pieceData)
	if bytes.Compare(hash.Sum(nil), piece.sha1[:]) != 0 {
		return errors.New("received corrupt piece")
	}

	piece.log.Debug("piece written successfully")
	piece.ok = true
	return nil
}

type peerHave struct {
	peer  *peerConn
	piece *piece
}

type peerBlock struct {
	peer  *peerConn
	block *block
	data  []byte
}
<commit_msg>do not send empty bitfield message<commit_after>package rain

import (
	"bytes"
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"sync"
	"time"

	"github.com/cenkalti/rain/internal/bitfield"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/protocol"
)

// All current implementations use 2^14 (16 kiB), and close connections which request an amount greater than that.
const blockSize = 16 * 1024

type peerConn struct {
	conn         net.Conn
	disconnected chan struct{} // will be closed when peer disconnects

	unchokeM       sync.Mutex    // protects unchokeC
	unchokeC       chan struct{} // will be closed when and "unchoke" message is received
	onceInterested sync.Once     // for sending "interested" message only once

	amChoking      bool // this client is choking the peer
	amInterested   bool // this client is interested in the peer
	peerChoking    bool // peer is choking this client
	peerInterested bool // peer is interested in this client
	// peerRequests   map[uint64]bool      // What remote peer requested
	// ourRequests    map[uint64]time.Time // What we requested, when we requested it

	log logger.Logger
}

func newPeerConn(conn net.Conn) *peerConn {
	return &peerConn{
		conn:         conn,
		disconnected: make(chan struct{}),
		unchokeC:     make(chan struct{}),
		amChoking:    true,
		peerChoking:  true,
		log:          logger.New("peer " + conn.RemoteAddr().String()),
	}
}

const connReadTimeout = 3 * time.Minute

// run processes incoming messages after handshake.
func (p *peerConn) run(t *transfer) {
	defer close(p.disconnected)
	p.log.Debugln("Communicating peer", p.conn.RemoteAddr())

	// Do not send bitfield if we don't have any pieces.
	// uTorrent seems to be dropping connections that send an empty bitfield message.
	if t.bitField.Count() != 0 {
		err := p.sendBitField(t.bitField)
		if err != nil {
			p.log.Error(err)
			return
		}
	}

	bitField := bitfield.New(nil, t.bitField.Len())

	first := true
	for {
		err := p.conn.SetReadDeadline(time.Now().Add(connReadTimeout))
		if err != nil {
			p.log.Error(err)
			return
		}

		var length uint32
		err = binary.Read(p.conn, binary.BigEndian, &length)
		if err != nil {
			if err == io.EOF {
				p.log.Warning("Remote peer has closed the connection")
				return
			}
			p.log.Error(err)
			return
		}
		p.log.Debugf("Received message of length: %d", length)

		if length == 0 { // keep-alive message
			p.log.Debug("Received message of type \"keep alive\"")
			continue
		}

		var msgType protocol.MessageType
		err = binary.Read(p.conn, binary.BigEndian, &msgType)
		if err != nil {
			p.log.Error(err)
			return
		}
		length--

		p.log.Debugf("Received message of type %q", msgType)

		switch msgType {
		case protocol.Choke:
			p.peerChoking = true
		case protocol.Unchoke:
			p.unchokeM.Lock()
			p.peerChoking = false
			close(p.unchokeC)
			// p.unchokeC = make(chan struct{})
			p.unchokeM.Unlock()
		case protocol.Interested:
			p.peerInterested = true
		case protocol.NotInterested:
			p.peerInterested = false
		case protocol.Have:
			var i uint32
			err = binary.Read(p.conn, binary.BigEndian, &i)
			if err != nil {
				p.log.Error(err)
				return
			}
			if i >= uint32(len(t.pieces)) {
				p.log.Error("unexpected piece index")
				return
			}
			piece := t.pieces[i]
			bitField.Set(i)
			p.log.Debug("Peer ", p.conn.RemoteAddr(), " has piece #", i)
			p.log.Debugln("new bitfield:", bitField.Hex())

			t.haveC <- peerHave{p, piece}
		case protocol.Bitfield:
			if !first {
				p.log.Error("bitfield can only be sent after handshake")
				return
			}

			if int64(length) != int64(len(bitField.Bytes())) {
				p.log.Error("invalid bitfield length")
				return
			}

			_, err = p.conn.Read(bitField.Bytes())
			if err != nil {
				p.log.Error(err)
				return
			}
			p.log.Debugln("Received bitfield:", bitField.Hex())

			for i := uint32(0); i < bitField.Len(); i++ {
				if bitField.Test(i) {
					t.haveC <- peerHave{p, t.pieces[i]}
				}
			}
		case protocol.Request:
		case protocol.Piece:
			var index uint32
			err = binary.Read(p.conn, binary.BigEndian, &index)
			if err != nil {
				p.log.Error(err)
				return
			}
			if index >= uint32(len(t.pieces)) {
				p.log.Error("unexpected piece index")
				return
			}
			piece := t.pieces[index]
			var begin uint32
			err = binary.Read(p.conn, binary.BigEndian, &begin)
			if err != nil {
				p.log.Error(err)
				return
			}
			if begin%blockSize != 0 {
				p.log.Error("unexpected piece offset")
				return
			}
			blockIndex := begin / blockSize
			if blockIndex >= uint32(len(piece.blocks)) {
				p.log.Error("unexpected piece offset")
				return
			}
			block := &piece.blocks[blockIndex]
			length -= 8
			if length != block.length {
				p.log.Error("unexpected block size")
				return
			}
			data := make([]byte, length)
			_, err = io.ReadFull(p.conn, data)
			if err != nil {
				p.log.Error(err)
				return
			}
			piece.blockC <- peerBlock{p, block, data}
		case protocol.Cancel:
		case protocol.Port:
		default:
			p.log.Debugf("Unknown message type: %d", msgType)
			// Discard remaining bytes.
			io.CopyN(ioutil.Discard, p.conn, int64(length))
		}

		first = false
	}
}

func (p *peerConn) sendBitField(b bitfield.BitField) error {
	var buf bytes.Buffer
	length := int32(1 + len(b.Bytes()))
	buf.Grow(4 + int(length))
	err := binary.Write(&buf, binary.BigEndian, length)
	if err != nil {
		return err
	}
	if err = buf.WriteByte(byte(protocol.Bitfield)); err != nil {
		return err
	}
	if _, err = buf.Write(b.Bytes()); err != nil {
		return err
	}
	p.log.Debugf("Sending message: \"bitfield\" %#v", buf.Bytes())
	_, err = io.Copy(p.conn, &buf)
	return err
}

// beInterested sends "interested" message to peer (once) and
// returns a channel that will be closed when an "unchoke" message is received.
func (p *peerConn) beInterested() (unchokeC chan struct{}, err error) {
	p.log.Debug("beInterested")
	p.unchokeM.Lock()
	defer p.unchokeM.Unlock()

	unchokeC = p.unchokeC

	if !p.peerChoking {
		return
	}

	p.onceInterested.Do(func() { err = p.sendMessage(protocol.Interested) })
	return
}

func (p *peerConn) sendMessage(msgType protocol.MessageType) error {
	var msg = struct {
		Length      uint32
		MessageType protocol.MessageType
	}{1, msgType}
	p.log.Debugf("Sending message: %q", msgType)
	return binary.Write(p.conn, binary.BigEndian, &msg)
}

type peerRequestMessage struct {
	ID                   protocol.MessageType
	Index, Begin, Length uint32
}

func newPeerRequestMessage(index, begin, length uint32) *peerRequestMessage {
	return &peerRequestMessage{protocol.Request, index, begin, length}
}

func (p *peerConn) sendRequest(m *peerRequestMessage) error {
	var msg = struct {
		Length  uint32
		Message peerRequestMessage
	}{13, *m}
	p.log.Debugf("Sending message: %q %#v", "request", msg)
	return binary.Write(p.conn, binary.BigEndian, &msg)
}

func (p *peerConn) downloadPiece(piece *piece) error {
	p.log.Debugf("downloading piece #%d", piece.index)

	unchokeC, err := p.beInterested()
	if err != nil {
		return err
	}

	select {
	case <-unchokeC:
	case <-time.After(time.Minute):
		p.conn.Close()
		return errors.New("Peer did not unchoke")
	}

	for _, b := range piece.blocks {
		if err := p.sendRequest(newPeerRequestMessage(piece.index, b.index*blockSize, b.length)); err != nil {
			return err
		}
	}

	pieceData := make([]byte, piece.length)
	for _ = range piece.blocks {
		select {
		case peerBlock := <-piece.blockC:
			p.log.Infoln("received block of length", len(peerBlock.data))
			copy(pieceData[peerBlock.block.index*blockSize:], peerBlock.data)
			if _, err = peerBlock.block.files.Write(peerBlock.data); err != nil {
				return err
			}
			piece.bitField.Set(peerBlock.block.index)
		case <-time.After(time.Minute):
			return fmt.Errorf("peer did not send piece #%d completely", piece.index)
		}
	}

	// Verify piece hash
	hash := sha1.New()
	hash.Write(pieceData)
	if bytes.Compare(hash.Sum(nil), piece.sha1[:]) != 0 {
		return errors.New("received corrupt piece")
	}

	piece.log.Debug("piece written successfully")
	piece.ok = true
	return nil
}

type peerHave struct {
	peer  *peerConn
	piece *piece
}

type peerBlock struct {
	peer  *peerConn
	block *block
	data  []byte
}
<|endoftext|>


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

package main

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

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

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

	pieces           map[uint32]struct{}
	queue            []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool
	closed           bool

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

	queueLock sync.Mutex

	id string
}

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

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

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

	peer.remoteChoked = true
	peer.localChoked = true
	return &peer
}

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

func (peer *Peer) close() {
	if !peer.closed {
		peer.connection.Close()
		peer.closed = true
	}
}

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

	var err error
	if peer.connection, err = net.DialTimeout("tcp", addr, time.Second*5); err != nil {
		return
	}
	defer peer.close()

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

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

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan *BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})
	peer.pieces = make(map[uint32]struct{})

	peer.torrent.addPeerChannel <- peer

	errorChannel := make(chan error)
	go peer.receiver(errorChannel)

	running := true
	for running {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case err := <-errorChannel:
			log.Printf("error in peer %s: %s", addr, err)
			peer.close()
			go func() {
				peer.torrent.removePeerChannel <- peer
			}()
			running = false
		case <-peer.done:
			peer.close()
		}
	}

	for {
		select {
		case <-peer.sendHaveChannel:
		case <-peer.sendPieceBlockChannel:
		case <-peer.requestPieceChannel:
		case <-peer.done:
			return
		}
	}
}

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

		length := binary.BigEndian.Uint32(lengthHeader)
		if length == 0 {
			// keep-alive
			continue
		}

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

		if err := peer.processMessage(length, data[0], data[1:]); err != nil {
			errorChannel <- err
			break
		}
	}
}

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

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

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

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

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

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

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

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

		index := binary.BigEndian.Uint32(payload)
		begin := binary.BigEndian.Uint32(payload[4:])
		blockLength := binary.BigEndian.Uint32(payload[8:])
		if blockLength > 32768 {
			return errors.New("peer requested length over 32KB")
		}

		peer.torrent.blockRequestChannel <- &BlockRequestMessage{peer, index, begin, blockLength}
	case PieceBlock:
		if length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		blockLength := length - 9
		if blockLength > 16384 {
			return errors.New("received block over 16KB")
		}

		index := binary.BigEndian.Uint32(payload)

		peer.queueLock.Lock()
		defer peer.queueLock.Unlock()

		idx := 0
		var piece *PeerPiece
		for k, v := range peer.queue {
			if v.index == index {
				idx = k
				piece = v
				break
			}
		}

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

		begin := binary.BigEndian.Uint32(payload[4:])

		blockIndex := begin / 16384
		if int(blockIndex) >= len(piece.blocks) {
			return errors.New("received too big block index")
		}
		piece.blocks[blockIndex] = payload[8:]

		piece.writes++
		if piece.writes == piece.reqWrites {
			// Glue all blocks together into a piece
			pieceData := []byte{}
			for k := range piece.blocks {
				pieceData = append(pieceData, piece.blocks[k]...)
			}

			// Send piece to Torrent
			peer.torrent.pieceChannel <- &PieceMessage{peer, index, pieceData}

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

		// TODO: Handle cancel
		/*
		   index := binary.BigEndian.Uint32(payload)
		   begin := binary.BigEndian.Uint32(payload[4:])
		   length := binary.BigEndian.Uint32(payload[8:])
		*/
	case Port:
		if length != 3 {
			return errors.New("length of port packet must be 3")
		}

		// port := binary.BigEndian.Uint16(payload)
		// Peer has a DHT node on port
	}
	return nil
}

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

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

func (peer *Peer) sendRequest(index, begin, length uint32) {
	packet := make([]byte, 17)
	binary.BigEndian.PutUint32(packet, 13) // Length
	packet[4] = Request
	binary.BigEndian.PutUint32(packet[5:], index)
	binary.BigEndian.PutUint32(packet[9:], begin)
	binary.BigEndian.PutUint32(packet[13:], length)
	peer.connection.Write(packet)
}

func (peer *Peer) sendPieceRequest(index uint32) {
	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(index)
	blocks := int(math.Ceil(float64(pieceLength) / 16384))
	piece := &PeerPiece{index, make([][]byte, blocks), 0, blocks}
	peer.queueLock.Lock()
	peer.queue = append(peer.queue, piece)
	peer.queueLock.Unlock()
	if !peer.remoteChoked {
		peer.requestPiece(piece)
	}
}

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

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	packet := make([]byte, 13)
	binary.BigEndian.PutUint32(packet, uint32(9+len(blockMessage.block)))
	packet[4] = PieceBlock
	binary.BigEndian.PutUint32(packet[5:], blockMessage.index)
	binary.BigEndian.PutUint32(packet[9:], blockMessage.begin)
	peer.connection.Write(packet)
	peer.connection.Write(blockMessage.block)
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := make([]byte, 9)
	packet[3] = 5 // Length
	packet[4] = Have
	binary.BigEndian.PutUint32(packet[5:], pieceIndex)
	peer.connection.Write(packet)
}

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

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

package main

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

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

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

	pieces           map[uint32]struct{}
	queue            []*PeerPiece
	connection       net.Conn
	remoteChoked     bool
	remoteInterested bool
	localChoked      bool
	localInterested  bool
	closed           bool

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

	queueLock sync.Mutex

	id string
}

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

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

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

	peer.remoteChoked = true
	peer.localChoked = true
	return &peer
}

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

func (peer *Peer) close() {
	if !peer.closed {
		peer.connection.Close()
		peer.closed = true
	}
}

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

	var err error
	if peer.connection, err = net.DialTimeout("tcp", addr, time.Second*5); err != nil {
		return
	}
	defer peer.close()

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

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

	peer.requestPieceChannel = make(chan uint32)
	peer.sendPieceBlockChannel = make(chan *BlockMessage)
	peer.sendHaveChannel = make(chan uint32)
	peer.done = make(chan struct{})
	peer.pieces = make(map[uint32]struct{})

	peer.torrent.addPeerChannel <- peer

	errorChannel := make(chan struct{})
	go peer.receiver(errorChannel)

	for {
		select {
		case pieceIndex := <-peer.requestPieceChannel:
			peer.sendPieceRequest(pieceIndex)
		case blockMessage := <-peer.sendPieceBlockChannel:
			peer.sendPieceBlockMessage(blockMessage)
		case pieceIndex := <-peer.sendHaveChannel:
			peer.sendHaveMessage(pieceIndex)
		case <-errorChannel:
			peer.close()
			go func() {
				peer.torrent.removePeerChannel <- peer
			}()
			for {
				select {
				// Until we receive another message over "done" to
				// confirm the removal of the peer, it's possible that
				// Torrent can send messages over other channels.
				case <-peer.sendHaveChannel:
				case <-peer.sendPieceBlockChannel:
				case <-peer.requestPieceChannel:
				case <-peer.done:
					return
				}
			}
		case <-peer.done:
			peer.close()
		}
	}
}

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

		length := binary.BigEndian.Uint32(lengthHeader)
		if length == 0 {
			// keep-alive
			continue
		}

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

		if err := peer.processMessage(length, data[0], data[1:]); err != nil {
			log.Print(err)
			errorChannel <- struct{}{}
			break
		}
	}
}

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

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

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

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

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

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

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

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

		index := binary.BigEndian.Uint32(payload)
		begin := binary.BigEndian.Uint32(payload[4:])
		blockLength := binary.BigEndian.Uint32(payload[8:])
		if blockLength > 32768 {
			return errors.New("peer requested length over 32KB")
		}

		peer.torrent.blockRequestChannel <- &BlockRequestMessage{peer, index, begin, blockLength}
	case PieceBlock:
		if length < 10 {
			return errors.New("length of piece packet must be at least 10")
		}

		blockLength := length - 9
		if blockLength > 16384 {
			return errors.New("received block over 16KB")
		}

		index := binary.BigEndian.Uint32(payload)

		peer.queueLock.Lock()
		defer peer.queueLock.Unlock()

		idx := 0
		var piece *PeerPiece
		for k, v := range peer.queue {
			if v.index == index {
				idx = k
				piece = v
				break
			}
		}

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

		begin := binary.BigEndian.Uint32(payload[4:])

		blockIndex := begin / 16384
		if int(blockIndex) >= len(piece.blocks) {
			return errors.New("received too big block index")
		}
		piece.blocks[blockIndex] = payload[8:]

		piece.writes++
		if piece.writes == piece.reqWrites {
			// Glue all blocks together into a piece
			pieceData := []byte{}
			for k := range piece.blocks {
				pieceData = append(pieceData, piece.blocks[k]...)
			}

			// Send piece to Torrent
			peer.torrent.pieceChannel <- &PieceMessage{peer, index, pieceData}

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

		// TODO: Handle cancel
		/*
		   index := binary.BigEndian.Uint32(payload)
		   begin := binary.BigEndian.Uint32(payload[4:])
		   length := binary.BigEndian.Uint32(payload[8:])
		*/
	case Port:
		if length != 3 {
			return errors.New("length of port packet must be 3")
		}

		// port := binary.BigEndian.Uint16(payload)
		// Peer has a DHT node on port
	}
	return nil
}

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

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

func (peer *Peer) sendRequest(index, begin, length uint32) {
	packet := make([]byte, 17)
	binary.BigEndian.PutUint32(packet, 13) // Length
	packet[4] = Request
	binary.BigEndian.PutUint32(packet[5:], index)
	binary.BigEndian.PutUint32(packet[9:], begin)
	binary.BigEndian.PutUint32(packet[13:], length)
	peer.connection.Write(packet)
}

func (peer *Peer) sendPieceRequest(index uint32) {
	peer.sendInterested()

	pieceLength := peer.torrent.getPieceLength(index)
	blocks := int(math.Ceil(float64(pieceLength) / 16384))
	piece := &PeerPiece{index, make([][]byte, blocks), 0, blocks}
	peer.queueLock.Lock()
	peer.queue = append(peer.queue, piece)
	peer.queueLock.Unlock()
	if !peer.remoteChoked {
		peer.requestPiece(piece)
	}
}

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

func (peer *Peer) sendPieceBlockMessage(blockMessage *BlockMessage) {
	packet := make([]byte, 13)
	binary.BigEndian.PutUint32(packet, uint32(9+len(blockMessage.block)))
	packet[4] = PieceBlock
	binary.BigEndian.PutUint32(packet[5:], blockMessage.index)
	binary.BigEndian.PutUint32(packet[9:], blockMessage.begin)
	peer.connection.Write(packet)
	peer.connection.Write(blockMessage.block)
}

func (peer *Peer) sendHaveMessage(pieceIndex uint32) {
	packet := make([]byte, 9)
	packet[3] = 5 // Length
	packet[4] = Have
	binary.BigEndian.PutUint32(packet[5:], pieceIndex)
	peer.connection.Write(packet)
}

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

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


<commit_before>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux darwin

package suid

import (
	"encoding/binary"
	"log"
	"os"
	"path/filepath"
	"syscall"

	"v.io/v23/verror"
)

var (
	errChownFailed        = verror.Register(pkgPath+".errChownFailed", verror.NoRetry, "{1:}{2:} os.Chown({3}, {4}, {5}) failed{:_}")
	errGetwdFailed        = verror.Register(pkgPath+".errGetwdFailed", verror.NoRetry, "{1:}{2:} os.Getwd failed{:_}")
	errStartProcessFailed = verror.Register(pkgPath+".errStartProcessFailed", verror.NoRetry, "{1:}{2:} syscall.StartProcess({3}) failed{:_}")
	errRemoveAllFailed    = verror.Register(pkgPath+".errRemoveAllFailed", verror.NoRetry, "{1:}{2:} os.RemoveAll({3}) failed{:_}")
	errFindProcessFailed  = verror.Register(pkgPath+".errFindProcessFailed", verror.NoRetry, "{1:}{2:} os.FindProcess({3}) failed{:_}")
	errKillFailed         = verror.Register(pkgPath+".errKillFailed", verror.NoRetry, "{1:}{2:} os.Process.Kill({3}) failed{:_}")
)

// Chown is only availabe on UNIX platforms so this file has a build
// restriction.
func (hw *WorkParameters) Chown() error {
	chown := func(path string, _ os.FileInfo, inerr error) error {
		if inerr != nil {
			return inerr
		}
		if hw.dryrun {
			log.Printf("[dryrun] os.Chown(%s, %d, %d)", path, hw.uid, hw.gid)
			return nil
		}
		return os.Chown(path, hw.uid, hw.gid)
	}

	chownPaths := hw.argv
	if !hw.chown {
		// Chown was invoked as part of regular suid execution, rather than directly
		// via --chown. In that case, we chown the workspace and log directory
		// TODO(rjkroege): Ensure that the device manager can read log entries.
		chownPaths = []string{hw.workspace, hw.logDir}
	}

	for _, p := range chownPaths {
		if err := filepath.Walk(p, chown); err != nil {
			return verror.New(errChownFailed, nil, p, hw.uid, hw.gid, err)
		}
	}
	return nil
}

func (hw *WorkParameters) Exec() error {
	attr := new(syscall.ProcAttr)

	dir, err := os.Getwd()
	if err != nil {
		log.Printf("error Getwd(): %v", err)
		return verror.New(errGetwdFailed, nil, err)
	}
	attr.Dir = dir
	attr.Env = hw.envv
	attr.Files = []uintptr{
		uintptr(syscall.Stdin),
		uintptr(syscall.Stdout),
		uintptr(syscall.Stderr),
	}

	attr.Sys = new(syscall.SysProcAttr)
	attr.Sys.Setsid = true
	if hw.dryrun {
		log.Printf("[dryrun] syscall.Setgid(%d)", hw.gid)
		log.Printf("[dryrun] syscall.Setuid(%d)", hw.uid)
	} else {
		attr.Sys.Credential = new(syscall.Credential)
		attr.Sys.Credential.Gid = uint32(hw.gid)
		attr.Sys.Credential.Uid = uint32(hw.uid)
	}

	// Make sure the child won't talk on the fd we use to talk back to the parent
	syscall.CloseOnExec(PipeToParentFD)

	// Start the child process
	pid, _, err := syscall.StartProcess(hw.argv0, hw.argv, attr)
	if err != nil {
		if !hw.dryrun {
			log.Printf("StartProcess failed: attr: %#v, attr.Sys: %#v, attr.Sys.Cred: %#v error: %v", attr, attr.Sys, attr.Sys.Credential, err)
		} else {
			log.Printf("StartProcess failed: %v", err)
		}
		return verror.New(errStartProcessFailed, nil, hw.argv0, err)
	}

	// Return the pid of the new child process
	pipeToParent := os.NewFile(PipeToParentFD, "pipe_to_parent_wr")
	if err = binary.Write(pipeToParent, binary.LittleEndian, int32(pid)); err != nil {
		log.Printf("Problem returning pid to parent: %v", err)
	} else {
		log.Printf("Returned pid %v to parent", pid)
	}

	os.Exit(0)
	return nil // Not reached.
}

func (hw *WorkParameters) Remove() error {
	for _, p := range hw.argv {
		if err := os.RemoveAll(p); err != nil {
			return verror.New(errRemoveAllFailed, nil, p, err)
		}
	}
	return nil
}

func (hw *WorkParameters) Kill() error {
	for _, pid := range hw.killPids {
		proc, err := os.FindProcess(pid)
		if err != nil {
			return verror.New(errFindProcessFailed, nil, pid, err)
		}

		if err = proc.Kill(); err != nil {
			return verror.New(errKillFailed, nil, pid, err)
		}
	}
	return nil
}
<commit_msg>services/device/internal/suid: don't complain if the app shutdown<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build linux darwin

package suid

import (
	"encoding/binary"
	"log"
	"os"
	"path/filepath"
	"syscall"

	"v.io/v23/verror"
)

var (
	errChownFailed        = verror.Register(pkgPath+".errChownFailed", verror.NoRetry, "{1:}{2:} os.Chown({3}, {4}, {5}) failed{:_}")
	errGetwdFailed        = verror.Register(pkgPath+".errGetwdFailed", verror.NoRetry, "{1:}{2:} os.Getwd failed{:_}")
	errStartProcessFailed = verror.Register(pkgPath+".errStartProcessFailed", verror.NoRetry, "{1:}{2:} syscall.StartProcess({3}) failed{:_}")
	errRemoveAllFailed    = verror.Register(pkgPath+".errRemoveAllFailed", verror.NoRetry, "{1:}{2:} os.RemoveAll({3}) failed{:_}")
	errFindProcessFailed  = verror.Register(pkgPath+".errFindProcessFailed", verror.NoRetry, "{1:}{2:} os.FindProcess({3}) failed{:_}")
	errKillFailed         = verror.Register(pkgPath+".errKillFailed", verror.NoRetry, "{1:}{2:} os.Process.Kill({3}) failed{:_}")
)

// Chown is only availabe on UNIX platforms so this file has a build
// restriction.
func (hw *WorkParameters) Chown() error {
	chown := func(path string, _ os.FileInfo, inerr error) error {
		if inerr != nil {
			return inerr
		}
		if hw.dryrun {
			log.Printf("[dryrun] os.Chown(%s, %d, %d)", path, hw.uid, hw.gid)
			return nil
		}
		return os.Chown(path, hw.uid, hw.gid)
	}

	chownPaths := hw.argv
	if !hw.chown {
		// Chown was invoked as part of regular suid execution, rather than directly
		// via --chown. In that case, we chown the workspace and log directory
		// TODO(rjkroege): Ensure that the device manager can read log entries.
		chownPaths = []string{hw.workspace, hw.logDir}
	}

	for _, p := range chownPaths {
		if err := filepath.Walk(p, chown); err != nil {
			return verror.New(errChownFailed, nil, p, hw.uid, hw.gid, err)
		}
	}
	return nil
}

func (hw *WorkParameters) Exec() error {
	attr := new(syscall.ProcAttr)

	dir, err := os.Getwd()
	if err != nil {
		log.Printf("error Getwd(): %v", err)
		return verror.New(errGetwdFailed, nil, err)
	}
	attr.Dir = dir
	attr.Env = hw.envv
	attr.Files = []uintptr{
		uintptr(syscall.Stdin),
		uintptr(syscall.Stdout),
		uintptr(syscall.Stderr),
	}

	attr.Sys = new(syscall.SysProcAttr)
	attr.Sys.Setsid = true
	if hw.dryrun {
		log.Printf("[dryrun] syscall.Setgid(%d)", hw.gid)
		log.Printf("[dryrun] syscall.Setuid(%d)", hw.uid)
	} else {
		attr.Sys.Credential = new(syscall.Credential)
		attr.Sys.Credential.Gid = uint32(hw.gid)
		attr.Sys.Credential.Uid = uint32(hw.uid)
	}

	// Make sure the child won't talk on the fd we use to talk back to the parent
	syscall.CloseOnExec(PipeToParentFD)

	// Start the child process
	pid, _, err := syscall.StartProcess(hw.argv0, hw.argv, attr)
	if err != nil {
		if !hw.dryrun {
			log.Printf("StartProcess failed: attr: %#v, attr.Sys: %#v, attr.Sys.Cred: %#v error: %v", attr, attr.Sys, attr.Sys.Credential, err)
		} else {
			log.Printf("StartProcess failed: %v", err)
		}
		return verror.New(errStartProcessFailed, nil, hw.argv0, err)
	}

	// Return the pid of the new child process
	pipeToParent := os.NewFile(PipeToParentFD, "pipe_to_parent_wr")
	if err = binary.Write(pipeToParent, binary.LittleEndian, int32(pid)); err != nil {
		log.Printf("Problem returning pid to parent: %v", err)
	} else {
		log.Printf("Returned pid %v to parent", pid)
	}

	os.Exit(0)
	return nil // Not reached.
}

func (hw *WorkParameters) Remove() error {
	for _, p := range hw.argv {
		if err := os.RemoveAll(p); err != nil {
			return verror.New(errRemoveAllFailed, nil, p, err)
		}
	}
	return nil
}

func (hw *WorkParameters) Kill() error {
	for _, pid := range hw.killPids {

		switch err := syscall.Kill(pid, 9); err {
		case syscall.ESRCH:
			// No such PID.
			log.Printf("process pid %d already killed", pid)
		default:
			// Something went wrong.
			return verror.New(errKillFailed, nil, pid, err)
		}

	}
	return nil
}
<|endoftext|>


<commit_before>package guber

// PodCollection is a Collection interface for Pods.
type PodCollection interface {
	Meta() *CollectionMeta
	New() *Pod
	Create(e *Pod) (*Pod, error)
	Query(q *QueryParams) (*PodList, error)
	List() (*PodList, error)
	Get(name string) (*Pod, error)
	Update(name string, r *Pod) (*Pod, error)
	Delete(name string) error
}

// Pods implmenets PodCollection.
type Pods struct {
	client    *RealClient
	Namespace string
}

// Meta implements the Collection interface.
func (c *Pods) Meta() *CollectionMeta {
	return &CollectionMeta{
		DomainName: "",
		APIGroup:   "api",
		APIVersion: "v1",
		APIName:    "pods",
		Kind:       "Pod",
	}
}

func (c *Pods) New() *Pod {
	return &Pod{
		collection: c,
	}
}

func (c *Pods) Create(e *Pod) (*Pod, error) {
	r := c.New()
	if err := c.client.Post().Collection(c).Namespace(c.Namespace).Entity(e).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Query(q *QueryParams) (*PodList, error) {
	list := new(PodList)
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Query(q).Do().Into(list); err != nil {
		return nil, err
	}
	for _, r := range list.Items {
		r.collection = c
	}
	return list, nil
}

func (c *Pods) List() (*PodList, error) {
	list := new(PodList)
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Do().Into(list); err != nil {
		return nil, err
	}
	for _, r := range list.Items {
		r.collection = c
	}
	return list, nil
}

func (c *Pods) Get(name string) (*Pod, error) {
	r := c.New()
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Name(name).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Update(name string, r *Pod) (*Pod, error) {
	if err := c.client.Patch().Collection(c).Namespace(c.Namespace).Name(name).Entity(r).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Delete(name string) error {
	req := c.client.Delete().Collection(c).Namespace(c.Namespace).Name(name).Do()
	return req.err
}

// Resource-level

func (r *Pod) Reload() (*Pod, error) {
	return r.collection.Get(r.Metadata.Name)
}

func (r *Pod) Save() error {
	_, err := r.collection.Update(r.Metadata.Name, r)
	return err
}

func (r *Pod) Delete() error {
	return r.collection.Delete(r.Metadata.Name)
}

func (r *Pod) Log(container string) (string, error) {
	// TODO we could consolidate all these collection-based methods with one Resource() mtehod
	return r.collection.client.Get().Collection(r.collection).Namespace(r.collection.Namespace).Name(r.Metadata.Name).Path("log?container=" + container).Do().Body()
}

func (r *Pod) IsReady() bool {
	if len(r.Status.Conditions) == 0 {
		return false
	}

	var readyCondition *PodStatusCondition
	for _, cond := range r.Status.Conditions {
		if cond.Type == "Ready" {
			readyCondition = cond
			break
		}
	}
	return readyCondition.Status == "True"
}
<commit_msg>Add pod.HeapsterStats()<commit_after>package guber

// PodCollection is a Collection interface for Pods.
type PodCollection interface {
	Meta() *CollectionMeta
	New() *Pod
	Create(e *Pod) (*Pod, error)
	Query(q *QueryParams) (*PodList, error)
	List() (*PodList, error)
	Get(name string) (*Pod, error)
	Update(name string, r *Pod) (*Pod, error)
	Delete(name string) error
}

// Pods implmenets PodCollection.
type Pods struct {
	client    *RealClient
	Namespace string
}

// Meta implements the Collection interface.
func (c *Pods) Meta() *CollectionMeta {
	return &CollectionMeta{
		DomainName: "",
		APIGroup:   "api",
		APIVersion: "v1",
		APIName:    "pods",
		Kind:       "Pod",
	}
}

func (c *Pods) New() *Pod {
	return &Pod{
		collection: c,
	}
}

func (c *Pods) Create(e *Pod) (*Pod, error) {
	r := c.New()
	if err := c.client.Post().Collection(c).Namespace(c.Namespace).Entity(e).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Query(q *QueryParams) (*PodList, error) {
	list := new(PodList)
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Query(q).Do().Into(list); err != nil {
		return nil, err
	}
	for _, r := range list.Items {
		r.collection = c
	}
	return list, nil
}

func (c *Pods) List() (*PodList, error) {
	list := new(PodList)
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Do().Into(list); err != nil {
		return nil, err
	}
	for _, r := range list.Items {
		r.collection = c
	}
	return list, nil
}

func (c *Pods) Get(name string) (*Pod, error) {
	r := c.New()
	if err := c.client.Get().Collection(c).Namespace(c.Namespace).Name(name).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Update(name string, r *Pod) (*Pod, error) {
	if err := c.client.Patch().Collection(c).Namespace(c.Namespace).Name(name).Entity(r).Do().Into(r); err != nil {
		return nil, err
	}
	return r, nil
}

func (c *Pods) Delete(name string) error {
	req := c.client.Delete().Collection(c).Namespace(c.Namespace).Name(name).Do()
	return req.err
}

// Resource-level

func (r *Pod) Reload() (*Pod, error) {
	return r.collection.Get(r.Metadata.Name)
}

func (r *Pod) Save() error {
	_, err := r.collection.Update(r.Metadata.Name, r)
	return err
}

func (r *Pod) Delete() error {
	return r.collection.Delete(r.Metadata.Name)
}

func (r *Pod) Log(container string) (string, error) {
	// TODO we could consolidate all these collection-based methods with one Resource() mtehod
	return r.collection.client.Get().Collection(r.collection).Namespace(r.collection.Namespace).Name(r.Metadata.Name).Path("log?container=" + container).Do().Body()
}

func (r *Pod) IsReady() bool {
	if len(r.Status.Conditions) == 0 {
		return false
	}

	var readyCondition *PodStatusCondition
	for _, cond := range r.Status.Conditions {
		if cond.Type == "Ready" {
			readyCondition = cond
			break
		}
	}
	return readyCondition.Status == "True"
}

func (r *Pod) HeapsterStats() (*HeapsterStats, error) {
	path := "api/v1/proxy/namespaces/kube-system/services/heapster/api/v1/model/namespaces/" + r.Metadata.Namespace + "/pods/" + r.Metadata.Name + "/stats"
	out := new(HeapsterStats)
	err := r.collection.client.Get().Path(path).Do().Into(out)
	return out, err
}
<|endoftext|>