<commit_before>package tips

// Reflection in Go:
// https://golang.org/pkg/reflect/

import (
	"reflect"
	"strings"
	"testing"
)

func TestReflectInt(t *testing.T) {
	i := 123
	ti := reflect.TypeOf(i)
	if ti.Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ti.Kind())
		return
	}
	v := reflect.ValueOf(i)
	if v.Type() != ti {
		t.Errorf("Expected %v but got %v", ti, v.Type())
		return
	}
	// Retrieve the actual value.
	var ri int64 = v.Int()
	if ri != int64(i) {
		t.Errorf("Expected %d but got %d", i, ri)
	}
	// Interestingly, this does not panic :)
	if v.String() != "<int Value>" {
		t.Errorf("Expected an empty string but got %q", v.String())
	}
	// This panics! (using unaddressable value)
	// v.SetInt(456)
}

func TestReflectIntPointer(t *testing.T) {
	i := 123
	var p interface{} = &i
	ty := reflect.TypeOf(p)
	if ty.Kind() != reflect.Ptr {
		t.Errorf("Expected reflect.Ptr but got %v", ty.Kind())
		return
	}
	v := reflect.ValueOf(p)
	if v.Type() != ty {
		t.Errorf("Expected %v but got %v", ty, v.Type())
		return
	}
	var ev reflect.Value = v.Elem()
	if ev.Type().Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ev.Type().Kind())
	}
	if ev.Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ev.Kind())
	}
	if !ev.CanSet() {
		t.Error("CanSet must return true.")
	}
	// Set value!
	v.Elem().SetInt(456)
	if i != 456 {
		t.Errorf("Expected value: %d", i)
	}
}

func TestReflectStruct(t *testing.T) {
	type entry struct {
		name string
		Age  int `myattr:"myvalue"`
	}
	e := entry{name: "Alice", Age: 13}

	ty := reflect.TypeOf(e)
	if ty.Kind() != reflect.Struct {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	if ty.NumField() != 2 {
		t.Errorf("Unexpected NumField(): %d", ty.NumField())
		return
	}
	var f0 reflect.StructField = ty.Field(0)
	if f0.Name != "name" || f0.Type.Kind() != reflect.String {
		t.Errorf("Unexpected field attrs: name = %q, type = %s", f0.Name, f0.Type)
	}
	f1 := ty.Field(1)
	if f1.Name != "Age" || f1.Type.Kind() != reflect.Int {
		t.Errorf("Unexpected field attrs: name = %q, type = %s", f1.Name, f1.Type)
	}
	// How to retrieve struct field tags.
	// https://golang.org/pkg/reflect/#StructTag
	if f1.Tag.Get("myattr") != "myvalue" {
		t.Errorf("Unexpected tag value: %q", f1.Tag.Get("myattr"))
	}

	// Value of struct.
	v := reflect.ValueOf(e)
	if v.Type() != ty {
		t.Errorf("Expected %v but got %v", ty, v.Type())
	}
	if v.NumField() != 2 {
		t.Errorf("Unexpected NumField(): %d", v.NumField())
		return
	}
	var vf0 reflect.Value = v.Field(0)
	if vf0.Type().Kind() != reflect.String {
		t.Errorf("Unexpected type for vf0: %v", vf0.Type())
	}
	if vf0.String() != "Alice" {
		t.Errorf("Unexpected value: %q", vf0.String())
	}
	vf1 := v.Field(1)
	if vf1.Type().Kind() != reflect.Int {
		t.Errorf("Unexpected type for vf1: %v", vf1.Type())
	}
	if vf0.String() != "Alice" {
		t.Errorf("Unexpected value: %q", vf0.String())
	}
	if vf1.Int() != 13 {
		t.Errorf("Unexpected value: %d", vf1.Int())
	}
	// fields of struct are not settable.
	if vf0.CanSet() || vf1.CanSet() {
		t.Error("CanSet of all fields must return false")
	}

	// This panics "using value obtained using unexported field".
	// vf0.SetString("hello")
	//
	// This panics with "using unaddressable value".
	// vf1.SetInt(20)
}

func TestReflectStructPointer(t *testing.T) {
	type entry struct {
		name string
		Age  int `myattr:"myvalue"`
	}
	p := &entry{name: "Alice", Age: 13}
	ty := reflect.TypeOf(p)
	if ty.Kind() != reflect.Ptr {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	// Set field values.
	v := reflect.ValueOf(p)
	v.Elem().Field(1).SetInt(24)
	if p.Age != 24 {
		t.Errorf("Unexpected value: %v", p.Age)
	}
	// This panics "using value obtained using unexported field".
	// v.Elem().Field(0).SetString("Bob")
}

func TestRelectEmbed(t *testing.T) {
	type embed struct {
		str string
		Str string
	}
	type Embed struct {
		Int     int
		integer int
	}
	type entry struct {
		F float32
		embed
		Embed
		float64
	}
	e := entry{
		F: 3.14,
		embed: embed{
			str: "hello",
			Str: "World",
		},
		Embed: Embed{
			Int:     100,
			integer: 10,
		},
		float64: 12.34,
	}
	ty := reflect.TypeOf(e)
	// Embeded struct are treated as fields as you may expect.
	if ty.NumField() != 4 {
		t.Errorf("Unexpected NumField() = %d", ty.NumField())
		return
	}
	f := ty.Field(0)
	if !(f.Name == "F" && !f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(1)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "embed" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(2)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "Embed" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(3)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "float64" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}

	// TODO: Investigate Value.
}

func TestReflectFunc(t *testing.T) {
	sum := func(x int, y int32) float32 {
		return float32(x + int(y))
	}
	var inter interface{} = sum
	ty := reflect.TypeOf(inter)
	if ty.Kind() != reflect.Func {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	// Types of arguments and return values.
	if ty.NumIn() != 2 {
		t.Errorf("Unexpected NumIn() = %v", ty.NumIn())
	}
	if ty.NumOut() != 1 {
		t.Errorf("Unexpected NumOut() = %v", ty.NumOut())
	}
	if ty.In(0).Kind() != reflect.Int {
		t.Errorf("Unexpected first arg type: %v", ty.In(0).Kind())
	}
	if ty.Out(0).Kind() != reflect.Float32 {
		t.Errorf("Unexpected first arg type: %v", ty.Out(0).Kind())
	}

	// Invoke
	v := reflect.ValueOf(inter)
	returns := v.Call([]reflect.Value{
		reflect.ValueOf(12),
		reflect.ValueOf(int32(34)),
	})
	if len(returns) != 1 {
		t.Errorf("Unexpected size of slice was returned: %d", len(returns))
		return
	}
	f := returns[0].Float()
	if f != 46.0 {
		t.Errorf("Unexpected return value: %f", f)
	}

	// Call panics if arguments are invalid.
	// v.Call([]reflect.Value{})
}

func TestReflectSlice(t *testing.T) {
	slicePtr := reflect.New(reflect.SliceOf(reflect.TypeOf(int(0))))
	for i := 0; i < 10; i++ {
		slicePtr.Elem().Set(reflect.Append(slicePtr.Elem(), reflect.ValueOf(i*i)))
	}
	s := slicePtr.Elem().Interface().([]int)
	expected := []int{0, 1, 4, 9, 16, 25, 36, 49, 64, 81}
	if !reflect.DeepEqual(s, expected) {
		t.Errorf("Expected %v but got %v", expected, s)
	}
}

func TestInterfaceToInterface(t *testing.T) {
	var err error
	var i interface{} = err
	if i != nil {
		t.Error("An interface from a nil interface must be nil too in Go.")
	}
	v := reflect.ValueOf(i)
	if v.IsValid() {
		t.Error("The value of nil must not be valid")
	}
	func() {
		defer func() {
			if recover() == nil {
				t.Error("v.IsNil() below must panic.")
			}
		}()
		// Wow, this panics!
		v.IsNil()
	}()
	func() {
		defer func() {
			p := recover()
			if p == nil {
				t.Error("The code must panic.")
				return
			}
			err := p.(error)
			if !strings.Contains(err.Error(), "interface is nil, not error") {
				t.Errorf("Unexpected error: %v", err)
			}
		}()
		err = i.(error)
	}()
}

func TestNilAndCall(t *testing.T) {
	errorType := reflect.TypeOf((*error)(nil)).Elem()
	if errorType.Name() != "error" {
		t.Errorf("Unexpected errorType: %v", errorType)
	}

	fn := func(e error) error {
		if e != nil {
			panic("Call this function with nil!")
		}
		return nil
	}
	v := reflect.ValueOf(fn)

	// NG
	// arg := reflect.ValueOf(nil)
	arg := reflect.New(errorType).Elem()
	rets := v.Call([]reflect.Value{arg})
	if len(rets) != 1 {
		t.Errorf("Unexpected len(ret): %d", len(rets))
		return
	}
	ret := rets[0]
	// Unlike the test above, a nil interface in the return values is not invalid and has a type.
	if !ret.IsValid() {
		t.Error("ret is invalid unexpectedly")
	}
	if !ret.IsNil() {
		t.Error("ret.IsNil() returned false unexpectedly")
	}
	if ret.Type() != errorType {
		t.Errorf("Unexpected type: %v", ret.Type())
	}
}

type myInt int

func TestReflectNamedType(t *testing.T) {
	var a interface{} = myInt(3)
	va := reflect.ValueOf(a)
	ta := reflect.TypeOf(a)
	intType := reflect.TypeOf(0)
	if ta.String() != "tips.myInt" {
		t.Errorf("ta.String() must be \"tips.myInt\", but got %q", ta.String())
	}
	if ta == intType {
		t.Errorf("typeof(myInt) == typeof(int) must be false")
	}
	if ta.Name() != "myInt" {
		t.Errorf("ta.Name() must be \"myInt\", but got %q", ta.Name())
	}
	if va.Kind() != reflect.Int {
		// Although typeof(myInt) != typeof(int), Kind() returns reflect.Int.
		t.Error("va.Kind() must be reflect.Int")
	}
	if va.Int() != 3 {
		t.Errorf("va.Int() must be 3, but got %v", va.Int())
	}
}

func TestReflectAddr(t *testing.T) {
	v := reflect.ValueOf(10)
	if v.CanAddr() {
		t.Errorf("v.CanAddr() must return false")
	}
	v = reflect.New(reflect.TypeOf(0))
	if !v.Elem().CanAddr() {
		t.Errorf("v.Elem().CanAddr() must return true")
	}
}
<commit_msg>Learn how to access unexported fields with reflect<commit_after>package tips

// Reflection in Go:
// https://golang.org/pkg/reflect/

import (
	"reflect"
	"strings"
	"testing"
	"unsafe"
)

func TestReflectInt(t *testing.T) {
	i := 123
	ti := reflect.TypeOf(i)
	if ti.Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ti.Kind())
		return
	}
	v := reflect.ValueOf(i)
	if v.Type() != ti {
		t.Errorf("Expected %v but got %v", ti, v.Type())
		return
	}
	// Retrieve the actual value.
	var ri int64 = v.Int()
	if ri != int64(i) {
		t.Errorf("Expected %d but got %d", i, ri)
	}
	// Interestingly, this does not panic :)
	if v.String() != "<int Value>" {
		t.Errorf("Expected an empty string but got %q", v.String())
	}
	// This panics! (using unaddressable value)
	// v.SetInt(456)
}

func TestReflectIntPointer(t *testing.T) {
	i := 123
	var p interface{} = &i
	ty := reflect.TypeOf(p)
	if ty.Kind() != reflect.Ptr {
		t.Errorf("Expected reflect.Ptr but got %v", ty.Kind())
		return
	}
	v := reflect.ValueOf(p)
	if v.Type() != ty {
		t.Errorf("Expected %v but got %v", ty, v.Type())
		return
	}
	var ev reflect.Value = v.Elem()
	if ev.Type().Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ev.Type().Kind())
	}
	if ev.Kind() != reflect.Int {
		t.Errorf("Expected reflect.Int but got %v", ev.Kind())
	}
	if !ev.CanSet() {
		t.Error("CanSet must return true.")
	}
	// Set value!
	v.Elem().SetInt(456)
	if i != 456 {
		t.Errorf("Expected value: %d", i)
	}
}

func TestReflectStruct(t *testing.T) {
	type entry struct {
		name string
		Age  int `myattr:"myvalue"`
	}
	e := entry{name: "Alice", Age: 13}

	ty := reflect.TypeOf(e)
	if ty.Kind() != reflect.Struct {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	if ty.NumField() != 2 {
		t.Errorf("Unexpected NumField(): %d", ty.NumField())
		return
	}
	var f0 reflect.StructField = ty.Field(0)
	if f0.Name != "name" || f0.Type.Kind() != reflect.String {
		t.Errorf("Unexpected field attrs: name = %q, type = %s", f0.Name, f0.Type)
	}
	f1 := ty.Field(1)
	if f1.Name != "Age" || f1.Type.Kind() != reflect.Int {
		t.Errorf("Unexpected field attrs: name = %q, type = %s", f1.Name, f1.Type)
	}
	// How to retrieve struct field tags.
	// https://golang.org/pkg/reflect/#StructTag
	if f1.Tag.Get("myattr") != "myvalue" {
		t.Errorf("Unexpected tag value: %q", f1.Tag.Get("myattr"))
	}

	// Value of struct.
	v := reflect.ValueOf(e)
	if v.Type() != ty {
		t.Errorf("Expected %v but got %v", ty, v.Type())
	}
	if v.NumField() != 2 {
		t.Errorf("Unexpected NumField(): %d", v.NumField())
		return
	}
	var vf0 reflect.Value = v.Field(0)
	if vf0.Type().Kind() != reflect.String {
		t.Errorf("Unexpected type for vf0: %v", vf0.Type())
	}
	if vf0.String() != "Alice" {
		t.Errorf("Unexpected value: %q", vf0.String())
	}
	vf1 := v.Field(1)
	if vf1.Type().Kind() != reflect.Int {
		t.Errorf("Unexpected type for vf1: %v", vf1.Type())
	}
	if vf0.String() != "Alice" {
		t.Errorf("Unexpected value: %q", vf0.String())
	}
	if vf1.Int() != 13 {
		t.Errorf("Unexpected value: %d", vf1.Int())
	}
	// fields of struct are not settable.
	if vf0.CanSet() || vf1.CanSet() {
		t.Error("CanSet of all fields must return false")
	}

	// This panics "using value obtained using unexported field".
	// vf0.SetString("hello")
	//
	// This panics with "using unaddressable value".
	// vf1.SetInt(20)
}

func TestReflectStructPointer(t *testing.T) {
	type entry struct {
		name string
		Age  int `myattr:"myvalue"`
	}
	p := &entry{name: "Alice", Age: 13}
	ty := reflect.TypeOf(p)
	if ty.Kind() != reflect.Ptr {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	// Set field values.
	v := reflect.ValueOf(p)
	v.Elem().Field(1).SetInt(24)
	if p.Age != 24 {
		t.Errorf("Unexpected value: %v", p.Age)
	}
	// This panics "using value obtained using unexported field".
	// v.Elem().Field(0).SetString("Bob")
}

func TestRelectEmbed(t *testing.T) {
	type embed struct {
		str string
		Str string
	}
	type Embed struct {
		Int     int
		integer int
	}
	type entry struct {
		F float32
		embed
		Embed
		float64
	}
	e := entry{
		F: 3.14,
		embed: embed{
			str: "hello",
			Str: "World",
		},
		Embed: Embed{
			Int:     100,
			integer: 10,
		},
		float64: 12.34,
	}
	ty := reflect.TypeOf(e)
	// Embeded struct are treated as fields as you may expect.
	if ty.NumField() != 4 {
		t.Errorf("Unexpected NumField() = %d", ty.NumField())
		return
	}
	f := ty.Field(0)
	if !(f.Name == "F" && !f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(1)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "embed" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(2)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "Embed" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}
	f = ty.Field(3)
	// Anonymous is true if the field is an embedded field.
	if !(f.Name == "float64" && f.Anonymous) {
		t.Errorf("Unexpected: %#v", f)
	}

	// TODO: Investigate Value.
}

func TestReflectFunc(t *testing.T) {
	sum := func(x int, y int32) float32 {
		return float32(x + int(y))
	}
	var inter interface{} = sum
	ty := reflect.TypeOf(inter)
	if ty.Kind() != reflect.Func {
		t.Errorf("Unexpected kind: %v", ty.Kind())
		return
	}
	// Types of arguments and return values.
	if ty.NumIn() != 2 {
		t.Errorf("Unexpected NumIn() = %v", ty.NumIn())
	}
	if ty.NumOut() != 1 {
		t.Errorf("Unexpected NumOut() = %v", ty.NumOut())
	}
	if ty.In(0).Kind() != reflect.Int {
		t.Errorf("Unexpected first arg type: %v", ty.In(0).Kind())
	}
	if ty.Out(0).Kind() != reflect.Float32 {
		t.Errorf("Unexpected first arg type: %v", ty.Out(0).Kind())
	}

	// Invoke
	v := reflect.ValueOf(inter)
	returns := v.Call([]reflect.Value{
		reflect.ValueOf(12),
		reflect.ValueOf(int32(34)),
	})
	if len(returns) != 1 {
		t.Errorf("Unexpected size of slice was returned: %d", len(returns))
		return
	}
	f := returns[0].Float()
	if f != 46.0 {
		t.Errorf("Unexpected return value: %f", f)
	}

	// Call panics if arguments are invalid.
	// v.Call([]reflect.Value{})
}

func TestReflectSlice(t *testing.T) {
	slicePtr := reflect.New(reflect.SliceOf(reflect.TypeOf(int(0))))
	for i := 0; i < 10; i++ {
		slicePtr.Elem().Set(reflect.Append(slicePtr.Elem(), reflect.ValueOf(i*i)))
	}
	s := slicePtr.Elem().Interface().([]int)
	expected := []int{0, 1, 4, 9, 16, 25, 36, 49, 64, 81}
	if !reflect.DeepEqual(s, expected) {
		t.Errorf("Expected %v but got %v", expected, s)
	}
}

func TestInterfaceToInterface(t *testing.T) {
	var err error
	var i interface{} = err
	if i != nil {
		t.Error("An interface from a nil interface must be nil too in Go.")
	}
	v := reflect.ValueOf(i)
	if v.IsValid() {
		t.Error("The value of nil must not be valid")
	}
	func() {
		defer func() {
			if recover() == nil {
				t.Error("v.IsNil() below must panic.")
			}
		}()
		// Wow, this panics!
		v.IsNil()
	}()
	func() {
		defer func() {
			p := recover()
			if p == nil {
				t.Error("The code must panic.")
				return
			}
			err := p.(error)
			if !strings.Contains(err.Error(), "interface is nil, not error") {
				t.Errorf("Unexpected error: %v", err)
			}
		}()
		err = i.(error)
	}()
}

func TestNilAndCall(t *testing.T) {
	errorType := reflect.TypeOf((*error)(nil)).Elem()
	if errorType.Name() != "error" {
		t.Errorf("Unexpected errorType: %v", errorType)
	}

	fn := func(e error) error {
		if e != nil {
			panic("Call this function with nil!")
		}
		return nil
	}
	v := reflect.ValueOf(fn)

	// NG
	// arg := reflect.ValueOf(nil)
	arg := reflect.New(errorType).Elem()
	rets := v.Call([]reflect.Value{arg})
	if len(rets) != 1 {
		t.Errorf("Unexpected len(ret): %d", len(rets))
		return
	}
	ret := rets[0]
	// Unlike the test above, a nil interface in the return values is not invalid and has a type.
	if !ret.IsValid() {
		t.Error("ret is invalid unexpectedly")
	}
	if !ret.IsNil() {
		t.Error("ret.IsNil() returned false unexpectedly")
	}
	if ret.Type() != errorType {
		t.Errorf("Unexpected type: %v", ret.Type())
	}
}

type myInt int

func TestReflectNamedType(t *testing.T) {
	var a interface{} = myInt(3)
	va := reflect.ValueOf(a)
	ta := reflect.TypeOf(a)
	intType := reflect.TypeOf(0)
	if ta.String() != "tips.myInt" {
		t.Errorf("ta.String() must be \"tips.myInt\", but got %q", ta.String())
	}
	if ta == intType {
		t.Errorf("typeof(myInt) == typeof(int) must be false")
	}
	if ta.Name() != "myInt" {
		t.Errorf("ta.Name() must be \"myInt\", but got %q", ta.Name())
	}
	if va.Kind() != reflect.Int {
		// Although typeof(myInt) != typeof(int), Kind() returns reflect.Int.
		t.Error("va.Kind() must be reflect.Int")
	}
	if va.Int() != 3 {
		t.Errorf("va.Int() must be 3, but got %v", va.Int())
	}
}

func TestReflectAddr(t *testing.T) {
	v := reflect.ValueOf(10)
	if v.CanAddr() {
		t.Errorf("v.CanAddr() must return false")
	}
	v = reflect.New(reflect.TypeOf(0))
	if !v.Elem().CanAddr() {
		t.Errorf("v.Elem().CanAddr() must return true")
	}
}

type childField struct {
	s string
}

type structWithUnexported struct {
	i     int8
	f     float32
	child childField
}

func expectUnexportedError(t *testing.T) {
	r := recover()
	if r == nil {
		t.Errorf("Expected panic but no panic was thrown")
	}
	if msg, _ := r.(string); !strings.Contains(msg, "unexported field") {
		t.Errorf("Unexpected error: %v", r)
	}
}

func TestStructUnexportedField(t *testing.T) {
	u := structWithUnexported{12, 3.4, childField{"hello"}}
	v := reflect.ValueOf(&u).Elem()

	// You can read unexported fields with reflect.
	if v.Field(0).Kind() != reflect.Int8 {
		t.Errorf("Expected %v but got %v", reflect.Int8, v.Field(0).Kind())
	}
	i := int8(v.Field(0).Int())
	if i != u.i {
		t.Errorf("Expected %v but got %v", u.i, i)
	}
	s := v.FieldByName("child").FieldByName("s").String()
	if s != u.child.s {
		t.Errorf("Expected %q but got %q", u.child.s, s)
	}

	// But you can not use Interface() to read unexported fields for some reasons.
	func() {
		defer expectUnexportedError(t)
		v.FieldByName("i").Interface()
	}()

	// You can not set unexported fields normally.
	func() {
		defer expectUnexportedError(t)
		v.FieldByName("f").SetFloat(5.6)
	}()
	// But, you can set values to unexported fields with unsafe.
	// Ref: https://stackoverflow.com/questions/42664837/access-unexported-fields-in-golang-reflect
	ff := v.FieldByName("f")
	newF := float32(7.8)
	reflect.NewAt(ff.Type(), unsafe.Pointer(ff.UnsafeAddr())).Elem().SetFloat(float64(newF))
	if u.f != newF {
		t.Errorf("Expected %v but got %v", newF, u.f)
	}
}
<|endoftext|>


<commit_before>package torrent

import (
	"bytes"
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/peerconn/peerreader"
	"github.com/cenkalti/rain/torrent/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/piecewriter"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/metainfo"
)

func (t *Torrent) handlePeerMessage(pm peer.Message) {
	pe := pm.Peer
	switch msg := pm.Message.(type) {
	case peerprotocol.HaveMessage:
		// Save have messages for processesing later received while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("unexpected piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		// pe.Logger().Debug("Peer ", pe.String(), " has piece #", pi.Index)
		t.piecePicker.HandleHave(pe, pi.Index)
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.BitfieldMessage:
		// Save bitfield messages while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		numBytes := uint32(bitfield.NumBytes(t.info.NumPieces))
		if uint32(len(msg.Data)) != numBytes {
			pe.Logger().Errorln("invalid bitfield length:", len(msg.Data))
			t.closePeer(pe)
			break
		}
		bf := bitfield.NewBytes(msg.Data, t.info.NumPieces)
		pe.Logger().Debugln("Received bitfield:", bf.Hex())
		for i := uint32(0); i < bf.Len(); i++ {
			if bf.Test(i) {
				t.piecePicker.HandleHave(pe, i)
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.HaveAllMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		for _, pi := range t.pieces {
			t.piecePicker.HandleHave(pe, pi.Index)
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.HaveNoneMessage: // TODO handle?
	case peerprotocol.AllowedFastMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid allowed fast piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		pe.Logger().Debug("Peer ", pe.String(), " has allowed fast for piece #", pi.Index)
		t.piecePicker.HandleAllowedFast(pe, msg.Index)
	case peerprotocol.UnchokeMessage:
		pe.PeerChoking = false
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.RequestBlocks(t.config.RequestQueueLength)
		}
		t.startPieceDownloaders()
	case peerprotocol.ChokeMessage:
		pe.PeerChoking = true
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.Choked()
			t.pieceDownloadersChoked[pe] = pd
			t.startPieceDownloaders()
		}
	case peerprotocol.InterestedMessage:
		// TODO handle intereseted messages
	case peerprotocol.NotInterestedMessage:
		// TODO handle not intereseted messages
	case peerreader.Piece:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("piece received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= uint32(len(t.pieces)) {
			pe.Logger().Errorln("invalid piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		piece := &t.pieces[msg.Index]
		block := piece.Blocks.Find(msg.Begin, uint32(len(msg.Data)))
		if block == nil {
			pe.Logger().Errorln("invalid piece begin:", msg.Begin, "length:", len(msg.Data))
			t.closePeer(pe)
			break
		}
		pe.BytesDownlaodedInChokePeriod += int64(len(msg.Data))
		t.bytesDownloaded += int64(len(msg.Data))
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			t.bytesWasted += int64(len(msg.Data))
			break
		}
		if pd.Piece.Index != msg.Index {
			t.bytesWasted += int64(len(msg.Data))
			break
		}
		pd.GotBlock(block, msg.Data)
		if !pd.Done() {
			pd.RequestBlocks(t.config.RequestQueueLength)
			pe.StartSnubTimer(t.config.RequestTimeout, t.peerSnubbedC)
			break
		}
		// t.log.Debugln("piece download completed. index:", pd.Piece.Index)
		t.closePieceDownloader(pd)
		pe.StopSnubTimer()

		ok = piece.VerifyHash(pd.Bytes, sha1.New()) // nolint: gosec
		if !ok {
			// TODO handle corrupt piece
			t.log.Error("received corrupt piece")
			t.closePeer(pd.Peer)
			t.startPieceDownloaders()
			break
		}

		for pe := range t.piecePicker.RequestedPeers(piece.Index) {
			pd2 := t.pieceDownloaders[pe]
			t.closePieceDownloader(pd2)
			pd2.CancelPending()
		}

		t.piecePicker.HandleWriting(piece.Index)
		pw := piecewriter.New(piece)
		t.pieceWriters[pw] = struct{}{}
		go pw.Run(pd.Bytes, t.pieceWriterResultC)

		t.startPieceDownloaders()
	case peerprotocol.RequestMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("request received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid request index:", msg.Index)
			t.closePeer(pe)
			break
		}
		if msg.Begin+msg.Length > t.pieces[msg.Index].Length {
			pe.Logger().Errorln("invalid request length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		if pe.AmChoking {
			if pe.FastExtension {
				m := peerprotocol.RejectMessage{RequestMessage: msg}
				pe.SendMessage(m)
			}
		} else {
			pe.SendPiece(msg, pi.Data)
		}
	case peerprotocol.RejectMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("reject received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid reject index:", msg.Index)
			t.closePeer(pe)
			break
		}
		piece := &t.pieces[msg.Index]
		block := piece.Blocks.Find(msg.Begin, msg.Length)
		if block == nil {
			pe.Logger().Errorln("invalid reject begin:", msg.Begin, "length:", msg.Length)
			t.closePeer(pe)
			break
		}
		// TODO check piece index
		// TODO check piece index on cancel
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			pe.Logger().Error("reject received but we don't have active download")
			t.closePeer(pe)
			break
		}
		pd.Rejected(block)
	case peerprotocol.CancelMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("cancel received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid cancel index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pe.CancelRequest(msg)
	case peerwriter.BlockUploaded:
		t.bytesUploaded += int64(msg.Length)
		pe.BytesUploadedInChokePeriod += int64(msg.Length)
	// TODO make extension messages value type
	case *peerprotocol.ExtensionHandshakeMessage:
		pe.Logger().Debugln("extension handshake received:", msg)
		if pe.ExtensionHandshake != nil {
			pe.Logger().Debugln("peer changed extensions")
			break
		}
		pe.ExtensionHandshake = msg

		if _, ok := msg.M[peerprotocol.ExtensionKeyMetadata]; ok {
			t.startInfoDownloaders()
		}
		if _, ok := msg.M[peerprotocol.ExtensionKeyPEX]; ok {
			if t.info != nil && t.info.Private != 1 {
				pe.StartPEX(t.peers)
			}
		}
	case *peerprotocol.ExtensionMetadataMessage:
		switch msg.Type {
		case peerprotocol.ExtensionMetadataMessageTypeRequest:
			extMsgID, ok := pe.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata]
			if !ok {
				break
			}
			if t.info == nil {
				// Send reject
				dataMsg := peerprotocol.ExtensionMetadataMessage{
					Type:  peerprotocol.ExtensionMetadataMessageTypeReject,
					Piece: msg.Piece,
				}
				extDataMsg := peerprotocol.ExtensionMessage{
					ExtendedMessageID: extMsgID,
					Payload:           &dataMsg,
				}
				pe.SendMessage(extDataMsg)
				break
			}
			if msg.Piece >= uint32(len(t.pieces)) {
				pe.Logger().Errorln("peer requested invalid metadata piece:", msg.Piece)
				t.closePeer(pe)
				break
			}
			// TODO Clients MAY implement flood protection by rejecting request messages after a certain number of them have been served. Typically the number of pieces of metadata times a factor.
			start := 16 * 1024 * msg.Piece
			end := 16 * 1024 * (msg.Piece + 1)
			totalSize := uint32(len(t.info.Bytes))
			if end > totalSize {
				end = totalSize
			}
			data := t.info.Bytes[start:end]
			dataMsg := peerprotocol.ExtensionMetadataMessage{
				Type:      peerprotocol.ExtensionMetadataMessageTypeData,
				Piece:     msg.Piece,
				TotalSize: totalSize,
				Data:      data,
			}
			extDataMsg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: extMsgID,
				Payload:           &dataMsg,
			}
			pe.SendMessage(extDataMsg)
		case peerprotocol.ExtensionMetadataMessageTypeData:
			id, ok := t.infoDownloaders[pe]
			if !ok {
				break
			}
			err := id.GotBlock(msg.Piece, msg.Data)
			if err != nil {
				pe.Logger().Error(err)
				t.closePeer(pe)
				t.startInfoDownloaders()
				break
			}
			if !id.Done() {
				id.RequestBlocks(t.config.RequestQueueLength)
				pe.StartSnubTimer(t.config.RequestTimeout, t.peerSnubbedC)
				break
			}
			pe.StopSnubTimer()

			hash := sha1.New()                              // nolint: gosec
			hash.Write(id.Bytes)                            // nolint: gosec
			if !bytes.Equal(hash.Sum(nil), t.infoHash[:]) { // nolint: gosec
				pe.Logger().Errorln("received info does not match with hash")
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
				break
			}
			t.stopInfoDownloaders()

			t.info, err = metainfo.NewInfo(id.Bytes)
			if err != nil {
				err = fmt.Errorf("cannot parse info bytes: %s", err)
				t.log.Error(err)
				t.stop(err)
				break
			}
			if t.info.Private == 1 {
				err = errors.New("private torrent from magnet")
				t.log.Error(err)
				t.stop(err)
				break
			}
			if t.resume != nil {
				err = t.resume.WriteInfo(t.info.Bytes)
				if err != nil {
					err = fmt.Errorf("cannot write resume info: %s", err)
					t.log.Error(err)
					t.stop(err)
					break
				}
			}
			t.startAllocator()
		case peerprotocol.ExtensionMetadataMessageTypeReject:
			id, ok := t.infoDownloaders[pe]
			if ok {
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
			}
		}
	case *peerprotocol.ExtensionPEXMessage:
		addrs, err := tracker.DecodePeersCompact([]byte(msg.Added))
		if err != nil {
			t.log.Error(err)
			break
		}
		t.handleNewPeers(addrs, "pex")
	default:
		panic(fmt.Sprintf("unhandled peer message type: %T", msg))
	}
}

func (t *Torrent) updateInterestedState(pe *peer.Peer) {
	if t.pieces == nil || t.bitfield == nil {
		return
	}
	interested := false
	for i := uint32(0); i < t.bitfield.Len(); i++ {
		weHave := t.bitfield.Test(i)
		peerHave := t.piecePicker.DoesHave(pe, i)
		if !weHave && peerHave {
			interested = true
			break
		}
	}
	if !pe.AmInterested && interested {
		pe.AmInterested = true
		msg := peerprotocol.InterestedMessage{}
		pe.SendMessage(msg)
		return
	}
	if pe.AmInterested && !interested {
		pe.AmInterested = false
		msg := peerprotocol.NotInterestedMessage{}
		pe.SendMessage(msg)
		return
	}
}
<commit_msg>remove todo<commit_after>package torrent

import (
	"bytes"
	"crypto/sha1" // nolint: gosec
	"errors"
	"fmt"

	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/peerconn/peerreader"
	"github.com/cenkalti/rain/torrent/internal/peerconn/peerwriter"
	"github.com/cenkalti/rain/torrent/internal/peerprotocol"
	"github.com/cenkalti/rain/torrent/internal/piecewriter"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/metainfo"
)

func (t *Torrent) handlePeerMessage(pm peer.Message) {
	pe := pm.Peer
	switch msg := pm.Message.(type) {
	case peerprotocol.HaveMessage:
		// Save have messages for processesing later received while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("unexpected piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		// pe.Logger().Debug("Peer ", pe.String(), " has piece #", pi.Index)
		t.piecePicker.HandleHave(pe, pi.Index)
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.BitfieldMessage:
		// Save bitfield messages while we don't have info yet.
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		numBytes := uint32(bitfield.NumBytes(t.info.NumPieces))
		if uint32(len(msg.Data)) != numBytes {
			pe.Logger().Errorln("invalid bitfield length:", len(msg.Data))
			t.closePeer(pe)
			break
		}
		bf := bitfield.NewBytes(msg.Data, t.info.NumPieces)
		pe.Logger().Debugln("Received bitfield:", bf.Hex())
		for i := uint32(0); i < bf.Len(); i++ {
			if bf.Test(i) {
				t.piecePicker.HandleHave(pe, i)
			}
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.HaveAllMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		for _, pi := range t.pieces {
			t.piecePicker.HandleHave(pe, pi.Index)
		}
		t.updateInterestedState(pe)
		t.startPieceDownloaders()
	case peerprotocol.HaveNoneMessage:
	case peerprotocol.AllowedFastMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Messages = append(pe.Messages, msg)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid allowed fast piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		pe.Logger().Debug("Peer ", pe.String(), " has allowed fast for piece #", pi.Index)
		t.piecePicker.HandleAllowedFast(pe, msg.Index)
	case peerprotocol.UnchokeMessage:
		pe.PeerChoking = false
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.RequestBlocks(t.config.RequestQueueLength)
		}
		t.startPieceDownloaders()
	case peerprotocol.ChokeMessage:
		pe.PeerChoking = true
		if pd, ok := t.pieceDownloaders[pe]; ok {
			pd.Choked()
			t.pieceDownloadersChoked[pe] = pd
			t.startPieceDownloaders()
		}
	case peerprotocol.InterestedMessage:
		// TODO handle intereseted messages
	case peerprotocol.NotInterestedMessage:
		// TODO handle not intereseted messages
	case peerreader.Piece:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("piece received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= uint32(len(t.pieces)) {
			pe.Logger().Errorln("invalid piece index:", msg.Index)
			t.closePeer(pe)
			break
		}
		piece := &t.pieces[msg.Index]
		block := piece.Blocks.Find(msg.Begin, uint32(len(msg.Data)))
		if block == nil {
			pe.Logger().Errorln("invalid piece begin:", msg.Begin, "length:", len(msg.Data))
			t.closePeer(pe)
			break
		}
		pe.BytesDownlaodedInChokePeriod += int64(len(msg.Data))
		t.bytesDownloaded += int64(len(msg.Data))
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			t.bytesWasted += int64(len(msg.Data))
			break
		}
		if pd.Piece.Index != msg.Index {
			t.bytesWasted += int64(len(msg.Data))
			break
		}
		pd.GotBlock(block, msg.Data)
		if !pd.Done() {
			pd.RequestBlocks(t.config.RequestQueueLength)
			pe.StartSnubTimer(t.config.RequestTimeout, t.peerSnubbedC)
			break
		}
		// t.log.Debugln("piece download completed. index:", pd.Piece.Index)
		t.closePieceDownloader(pd)
		pe.StopSnubTimer()

		ok = piece.VerifyHash(pd.Bytes, sha1.New()) // nolint: gosec
		if !ok {
			// TODO handle corrupt piece
			t.log.Error("received corrupt piece")
			t.closePeer(pd.Peer)
			t.startPieceDownloaders()
			break
		}

		for pe := range t.piecePicker.RequestedPeers(piece.Index) {
			pd2 := t.pieceDownloaders[pe]
			t.closePieceDownloader(pd2)
			pd2.CancelPending()
		}

		t.piecePicker.HandleWriting(piece.Index)
		pw := piecewriter.New(piece)
		t.pieceWriters[pw] = struct{}{}
		go pw.Run(pd.Bytes, t.pieceWriterResultC)

		t.startPieceDownloaders()
	case peerprotocol.RequestMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("request received but we don't have info")
			t.closePeer(pe)
			break
		}
		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid request index:", msg.Index)
			t.closePeer(pe)
			break
		}
		if msg.Begin+msg.Length > t.pieces[msg.Index].Length {
			pe.Logger().Errorln("invalid request length:", msg.Length)
			t.closePeer(pe)
			break
		}
		pi := &t.pieces[msg.Index]
		if pe.AmChoking {
			if pe.FastExtension {
				m := peerprotocol.RejectMessage{RequestMessage: msg}
				pe.SendMessage(m)
			}
		} else {
			pe.SendPiece(msg, pi.Data)
		}
	case peerprotocol.RejectMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("reject received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid reject index:", msg.Index)
			t.closePeer(pe)
			break
		}
		piece := &t.pieces[msg.Index]
		block := piece.Blocks.Find(msg.Begin, msg.Length)
		if block == nil {
			pe.Logger().Errorln("invalid reject begin:", msg.Begin, "length:", msg.Length)
			t.closePeer(pe)
			break
		}
		// TODO check piece index
		// TODO check piece index on cancel
		pd, ok := t.pieceDownloaders[pe]
		if !ok {
			pe.Logger().Error("reject received but we don't have active download")
			t.closePeer(pe)
			break
		}
		pd.Rejected(block)
	case peerprotocol.CancelMessage:
		if t.pieces == nil || t.bitfield == nil {
			pe.Logger().Error("cancel received but we don't have info")
			t.closePeer(pe)
			break
		}

		if msg.Index >= t.info.NumPieces {
			pe.Logger().Errorln("invalid cancel index:", msg.Index)
			t.closePeer(pe)
			break
		}
		pe.CancelRequest(msg)
	case peerwriter.BlockUploaded:
		t.bytesUploaded += int64(msg.Length)
		pe.BytesUploadedInChokePeriod += int64(msg.Length)
	// TODO make extension messages value type
	case *peerprotocol.ExtensionHandshakeMessage:
		pe.Logger().Debugln("extension handshake received:", msg)
		if pe.ExtensionHandshake != nil {
			pe.Logger().Debugln("peer changed extensions")
			break
		}
		pe.ExtensionHandshake = msg

		if _, ok := msg.M[peerprotocol.ExtensionKeyMetadata]; ok {
			t.startInfoDownloaders()
		}
		if _, ok := msg.M[peerprotocol.ExtensionKeyPEX]; ok {
			if t.info != nil && t.info.Private != 1 {
				pe.StartPEX(t.peers)
			}
		}
	case *peerprotocol.ExtensionMetadataMessage:
		switch msg.Type {
		case peerprotocol.ExtensionMetadataMessageTypeRequest:
			extMsgID, ok := pe.ExtensionHandshake.M[peerprotocol.ExtensionKeyMetadata]
			if !ok {
				break
			}
			if t.info == nil {
				// Send reject
				dataMsg := peerprotocol.ExtensionMetadataMessage{
					Type:  peerprotocol.ExtensionMetadataMessageTypeReject,
					Piece: msg.Piece,
				}
				extDataMsg := peerprotocol.ExtensionMessage{
					ExtendedMessageID: extMsgID,
					Payload:           &dataMsg,
				}
				pe.SendMessage(extDataMsg)
				break
			}
			if msg.Piece >= uint32(len(t.pieces)) {
				pe.Logger().Errorln("peer requested invalid metadata piece:", msg.Piece)
				t.closePeer(pe)
				break
			}
			// TODO Clients MAY implement flood protection by rejecting request messages after a certain number of them have been served. Typically the number of pieces of metadata times a factor.
			start := 16 * 1024 * msg.Piece
			end := 16 * 1024 * (msg.Piece + 1)
			totalSize := uint32(len(t.info.Bytes))
			if end > totalSize {
				end = totalSize
			}
			data := t.info.Bytes[start:end]
			dataMsg := peerprotocol.ExtensionMetadataMessage{
				Type:      peerprotocol.ExtensionMetadataMessageTypeData,
				Piece:     msg.Piece,
				TotalSize: totalSize,
				Data:      data,
			}
			extDataMsg := peerprotocol.ExtensionMessage{
				ExtendedMessageID: extMsgID,
				Payload:           &dataMsg,
			}
			pe.SendMessage(extDataMsg)
		case peerprotocol.ExtensionMetadataMessageTypeData:
			id, ok := t.infoDownloaders[pe]
			if !ok {
				break
			}
			err := id.GotBlock(msg.Piece, msg.Data)
			if err != nil {
				pe.Logger().Error(err)
				t.closePeer(pe)
				t.startInfoDownloaders()
				break
			}
			if !id.Done() {
				id.RequestBlocks(t.config.RequestQueueLength)
				pe.StartSnubTimer(t.config.RequestTimeout, t.peerSnubbedC)
				break
			}
			pe.StopSnubTimer()

			hash := sha1.New()                              // nolint: gosec
			hash.Write(id.Bytes)                            // nolint: gosec
			if !bytes.Equal(hash.Sum(nil), t.infoHash[:]) { // nolint: gosec
				pe.Logger().Errorln("received info does not match with hash")
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
				break
			}
			t.stopInfoDownloaders()

			t.info, err = metainfo.NewInfo(id.Bytes)
			if err != nil {
				err = fmt.Errorf("cannot parse info bytes: %s", err)
				t.log.Error(err)
				t.stop(err)
				break
			}
			if t.info.Private == 1 {
				err = errors.New("private torrent from magnet")
				t.log.Error(err)
				t.stop(err)
				break
			}
			if t.resume != nil {
				err = t.resume.WriteInfo(t.info.Bytes)
				if err != nil {
					err = fmt.Errorf("cannot write resume info: %s", err)
					t.log.Error(err)
					t.stop(err)
					break
				}
			}
			t.startAllocator()
		case peerprotocol.ExtensionMetadataMessageTypeReject:
			id, ok := t.infoDownloaders[pe]
			if ok {
				t.closePeer(id.Peer)
				t.startInfoDownloaders()
			}
		}
	case *peerprotocol.ExtensionPEXMessage:
		addrs, err := tracker.DecodePeersCompact([]byte(msg.Added))
		if err != nil {
			t.log.Error(err)
			break
		}
		t.handleNewPeers(addrs, "pex")
	default:
		panic(fmt.Sprintf("unhandled peer message type: %T", msg))
	}
}

func (t *Torrent) updateInterestedState(pe *peer.Peer) {
	if t.pieces == nil || t.bitfield == nil {
		return
	}
	interested := false
	for i := uint32(0); i < t.bitfield.Len(); i++ {
		weHave := t.bitfield.Test(i)
		peerHave := t.piecePicker.DoesHave(pe, i)
		if !weHave && peerHave {
			interested = true
			break
		}
	}
	if !pe.AmInterested && interested {
		pe.AmInterested = true
		msg := peerprotocol.InterestedMessage{}
		pe.SendMessage(msg)
		return
	}
	if pe.AmInterested && !interested {
		pe.AmInterested = false
		msg := peerprotocol.NotInterestedMessage{}
		pe.SendMessage(msg)
		return
	}
}
<|endoftext|>


<commit_before>package bitset

import (
	"bytes"
	"fmt"
	"math"
)

const (
	sw_64   uint64 = 64
	slg2_64 uint64 = 5
	m1_64   uint64 = 0x5555555555555555 // 0101...
	m2_64   uint64 = 0x3333333333333333 // 00110011..
	m4_64   uint64 = 0x0f0f0f0f0f0f0f0f // 00001111...
	hff_64  uint64 = 0xffffffffffffffff // all ones
)

func wordsNeeded64(n uint64) uint64 {
	if n == 0 {
		return 1
	} else if n > math.MaxUint64-sw_64 {
		return math.MaxUint64 >> slg2_64
	}
	return (n + (sw_64 - 1)) >> slg2_64
}

type Bitset64 struct {
	n uint64
	b []uint64
}

// Returns the current size of the bitset.
func (b *Bitset64) Len() uint64 {
	return b.n
}

// Test whether bit i is set.
func (b *Bitset64) Test(i uint64) bool {
	if i >= b.n {
		return false
	}
	return ((b.b[i>>slg2_64] & (1 << (i & (sw_64 - 1)))) != 0)
}

// Set bit i to 1.
func (b *Bitset64) Set(i uint64) {
	if i >= b.n {
		nsize := wordsNeeded64(i + 1)
		l := uint64(len(b.b))
		if nsize > l {
			nb := make([]uint64, nsize-l)
			b.b = append(b.b, nb...)
		}
		b.n = i + 1
	}
	b.b[i>>slg2_64] |= (1 << (i & (sw_64 - 1)))
}

// Set bit i to 0.
func (b *Bitset64) Clear(i uint64) {
	if i >= b.n {
		return
	}
	b.b[i>>slg2_64] &^= 1 << (i & (sw_64 - 1))
}

// Flip bit i.
func (b *Bitset64) Flip(i uint64) {
	if i >= b.n {
		b.Set(i)
	}
	b.b[i>>slg2_64] ^= 1 << (i & (sw_64 - 1))
}

// Clear all bits in the bitset.
func (b *Bitset64) Reset() {
	for i := range b.b {
		b.b[i] = 0
	}
}

// Get the number of words used in the bitset.
func (b *Bitset64) wordCount() uint64 {
	return wordsNeeded64(b.n)
}

// Clone the bitset.
func (b *Bitset64) Clone() *Bitset64 {
	c := New64(b.n)
	copy(c.b, b.b)
	return c
}

// Copy the bitset into another bitset, returning the size of the destination
// bitset.
func (b *Bitset64) Copy(c *Bitset64) (n uint64) {
	copy(c.b, b.b)
	n = c.n
	if b.n < c.n {
		n = b.n
	}
	return
}

func popCountUint64(x uint64) uint64 {
	x -= (x >> 1) & m1_64                // put count of each 2 bits into those 2 bits
	x = (x & m2_64) + ((x >> 2) & m2_64) // put count of each 4 bits into those 4 bits
	x = (x + (x >> 4)) & m4_64           // put count of each 8 bits into those 8 bits
	x += x >> 8                          // put count of each 16 bits into their lowest 8 bits
	x += x >> 16                         // put count of each 32 bits into their lowest 8 bits
	x += x >> 32                         // put count of each 64 bits into their lowest 8 bits
	return x & 0x7f
}

// Get the number of set bits in the bitset.
func (b *Bitset64) Count() uint64 {
	sum := uint64(0)
	for _, w := range b.b {
		sum += popCountUint64(w)
	}
	return sum
}

// Test if two bitsets are equal. Returns true if both bitsets are the same
// size and all the same bits are set in both bitsets.
func (b *Bitset64) Equal(c *Bitset64) bool {
	if b.n != c.n {
		return false
	}
	for p, v := range b.b {
		if c.b[p] != v {
			return false
		}
	}
	return true
}

// Bitset &^ (and or); difference between receiver and another set.
func (b *Bitset64) Difference(ob *Bitset64) (result *Bitset64) {
	result = b.Clone() // clone b (in case b is bigger than ob)
	szl := ob.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] &^ ob.b[i]
	}
	return
}

func sortByLength64(a *Bitset64, b *Bitset64) (ap *Bitset64, bp *Bitset64) {
	if a.n <= b.n {
		ap, bp = a, b
	} else {
		ap, bp = b, a
	}
	return
}

// Bitset & (and); intersection of receiver and another set.
func (b *Bitset64) Intersection(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	result = New64(b.n)
	for i, w := range b.b {
		result.b[i] = w & ob.b[i]
	}
	return
}

// Bitset | (or); union of receiver and another set.
func (b *Bitset64) Union(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	result = ob.Clone()
	szl := ob.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] | ob.b[i]
	}
	return
}

// Bitset ^ (xor); symmetric difference of receiver and another set.
func (b *Bitset64) SymmetricDifference(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	// ob is bigger, so clone it
	result = ob.Clone()
	szl := b.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] ^ ob.b[i]
	}
	return
}

// Return true if the bitset's length is a multiple of the word size.
func (b *Bitset64) isEven() bool {
	return (b.n % sw_64) == 0
}

// Clean last word by setting unused bits to 0.
func (b *Bitset64) cleanLastWord() {
	if !b.isEven() {
		b.b[wordsNeeded64(b.n)-1] &= (hff_64 >> (sw_64 - (b.n % sw_64)))
	}
}

// Return the (local) complement of a bitset (up to n bits).
// func (b *Bitset64) Complement() (result *Bitset64) {
// 	result = New64(b.n)
// 	for i, w := range b.b {
// 		result.b[i] = ^(w)
// 	}
// 	result.cleanLastWord()
// 	return
// }

// Returns true if all bits in the bitset are set.
func (b *Bitset64) All() bool {
	return b.Count() == b.n
}

// Returns true if no bit in the bitset is set.
func (b *Bitset64) None() bool {
	for _, w := range b.b {
		if w > 0 {
			return false
		}
	}
	return true
}

// Return true if any bit in the bitset is set.
func (b *Bitset64) Any() bool {
	return !b.None()
}

// Get a string representation of the words in the bitset.
func (b *Bitset64) String() string {
	f := bytes.NewBufferString("")
	for i := int(wordsNeeded64(b.n) - 1); i >= 0; i-- {
		fmt.Fprintf(f, "%064b.", b.b[i])
	}
	return f.String()
}

// Make a new bitset with a starting capacity of n bits. The bitset expands
// automatically.
func New64(n uint64) *Bitset64 {
	nWords := wordsNeeded64(n)
	if nWords > math.MaxInt32-1 {
		panic(fmt.Sprintf("Bitset64 needs %d %d-bit words to store %d bits, but slices cannot hold more than %d items.", nWords, sw_64, n, math.MaxInt32-1))
	}
	b := &Bitset64{
		n,
		make([]uint64, nWords),
	}
	return b
}
<commit_msg>log2(64) is 6<commit_after>package bitset

import (
	"bytes"
	"fmt"
	"math"
)

const (
	sw_64   uint64 = 64
	slg2_64 uint64 = 6
	m1_64   uint64 = 0x5555555555555555 // 0101...
	m2_64   uint64 = 0x3333333333333333 // 00110011..
	m4_64   uint64 = 0x0f0f0f0f0f0f0f0f // 00001111...
	hff_64  uint64 = 0xffffffffffffffff // all ones
)

func wordsNeeded64(n uint64) uint64 {
	if n == 0 {
		return 1
	} else if n > math.MaxUint64-sw_64 {
		return math.MaxUint64 >> slg2_64
	}
	return (n + (sw_64 - 1)) >> slg2_64
}

type Bitset64 struct {
	n uint64
	b []uint64
}

// Returns the current size of the bitset.
func (b *Bitset64) Len() uint64 {
	return b.n
}

// Test whether bit i is set.
func (b *Bitset64) Test(i uint64) bool {
	if i >= b.n {
		return false
	}
	return ((b.b[i>>slg2_64] & (1 << (i & (sw_64 - 1)))) != 0)
}

// Set bit i to 1.
func (b *Bitset64) Set(i uint64) {
	if i >= b.n {
		nsize := wordsNeeded64(i + 1)
		l := uint64(len(b.b))
		if nsize > l {
			nb := make([]uint64, nsize-l)
			b.b = append(b.b, nb...)
		}
		b.n = i + 1
	}
	b.b[i>>slg2_64] |= (1 << (i & (sw_64 - 1)))
}

// Set bit i to 0.
func (b *Bitset64) Clear(i uint64) {
	if i >= b.n {
		return
	}
	b.b[i>>slg2_64] &^= 1 << (i & (sw_64 - 1))
}

// Flip bit i.
func (b *Bitset64) Flip(i uint64) {
	if i >= b.n {
		b.Set(i)
	}
	b.b[i>>slg2_64] ^= 1 << (i & (sw_64 - 1))
}

// Clear all bits in the bitset.
func (b *Bitset64) Reset() {
	for i := range b.b {
		b.b[i] = 0
	}
}

// Get the number of words used in the bitset.
func (b *Bitset64) wordCount() uint64 {
	return wordsNeeded64(b.n)
}

// Clone the bitset.
func (b *Bitset64) Clone() *Bitset64 {
	c := New64(b.n)
	copy(c.b, b.b)
	return c
}

// Copy the bitset into another bitset, returning the size of the destination
// bitset.
func (b *Bitset64) Copy(c *Bitset64) (n uint64) {
	copy(c.b, b.b)
	n = c.n
	if b.n < c.n {
		n = b.n
	}
	return
}

func popCountUint64(x uint64) uint64 {
	x -= (x >> 1) & m1_64                // put count of each 2 bits into those 2 bits
	x = (x & m2_64) + ((x >> 2) & m2_64) // put count of each 4 bits into those 4 bits
	x = (x + (x >> 4)) & m4_64           // put count of each 8 bits into those 8 bits
	x += x >> 8                          // put count of each 16 bits into their lowest 8 bits
	x += x >> 16                         // put count of each 32 bits into their lowest 8 bits
	x += x >> 32                         // put count of each 64 bits into their lowest 8 bits
	return x & 0x7f
}

// Get the number of set bits in the bitset.
func (b *Bitset64) Count() uint64 {
	sum := uint64(0)
	for _, w := range b.b {
		sum += popCountUint64(w)
	}
	return sum
}

// Test if two bitsets are equal. Returns true if both bitsets are the same
// size and all the same bits are set in both bitsets.
func (b *Bitset64) Equal(c *Bitset64) bool {
	if b.n != c.n {
		return false
	}
	for p, v := range b.b {
		if c.b[p] != v {
			return false
		}
	}
	return true
}

// Bitset &^ (and or); difference between receiver and another set.
func (b *Bitset64) Difference(ob *Bitset64) (result *Bitset64) {
	result = b.Clone() // clone b (in case b is bigger than ob)
	szl := ob.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] &^ ob.b[i]
	}
	return
}

func sortByLength64(a *Bitset64, b *Bitset64) (ap *Bitset64, bp *Bitset64) {
	if a.n <= b.n {
		ap, bp = a, b
	} else {
		ap, bp = b, a
	}
	return
}

// Bitset & (and); intersection of receiver and another set.
func (b *Bitset64) Intersection(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	result = New64(b.n)
	for i, w := range b.b {
		result.b[i] = w & ob.b[i]
	}
	return
}

// Bitset | (or); union of receiver and another set.
func (b *Bitset64) Union(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	result = ob.Clone()
	szl := ob.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] | ob.b[i]
	}
	return
}

// Bitset ^ (xor); symmetric difference of receiver and another set.
func (b *Bitset64) SymmetricDifference(ob *Bitset64) (result *Bitset64) {
	b, ob = sortByLength64(b, ob)
	// ob is bigger, so clone it
	result = ob.Clone()
	szl := b.wordCount()
	l := uint64(len(b.b))
	for i := uint64(0); i < l; i++ {
		if i >= szl {
			break
		}
		result.b[i] = b.b[i] ^ ob.b[i]
	}
	return
}

// Return true if the bitset's length is a multiple of the word size.
func (b *Bitset64) isEven() bool {
	return (b.n % sw_64) == 0
}

// Clean last word by setting unused bits to 0.
func (b *Bitset64) cleanLastWord() {
	if !b.isEven() {
		b.b[wordsNeeded64(b.n)-1] &= (hff_64 >> (sw_64 - (b.n % sw_64)))
	}
}

// Return the (local) complement of a bitset (up to n bits).
// func (b *Bitset64) Complement() (result *Bitset64) {
// 	result = New64(b.n)
// 	for i, w := range b.b {
// 		result.b[i] = ^(w)
// 	}
// 	result.cleanLastWord()
// 	return
// }

// Returns true if all bits in the bitset are set.
func (b *Bitset64) All() bool {
	return b.Count() == b.n
}

// Returns true if no bit in the bitset is set.
func (b *Bitset64) None() bool {
	for _, w := range b.b {
		if w > 0 {
			return false
		}
	}
	return true
}

// Return true if any bit in the bitset is set.
func (b *Bitset64) Any() bool {
	return !b.None()
}

// Get a string representation of the words in the bitset.
func (b *Bitset64) String() string {
	f := bytes.NewBufferString("")
	for i := int(wordsNeeded64(b.n) - 1); i >= 0; i-- {
		fmt.Fprintf(f, "%064b.", b.b[i])
	}
	return f.String()
}

// Make a new bitset with a starting capacity of n bits. The bitset expands
// automatically.
func New64(n uint64) *Bitset64 {
	nWords := wordsNeeded64(n)
	if nWords > math.MaxInt32-1 {
		panic(fmt.Sprintf("Bitset64 needs %d %d-bit words to store %d bits, but slices cannot hold more than %d items.", nWords, sw_64, n, math.MaxInt32-1))
	}
	b := &Bitset64{
		n,
		make([]uint64, nWords),
	}
	return b
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"
	"log"
	"time"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
	"github.com/coreos/coreinit/registry"
)

const (
	DefaultJobWatchClaimTTL = "10s"
	DefaultRefreshInterval  = "60s"
)

type JobWatcher struct {
	registry        *registry.Registry
	scheduler       *Scheduler
	machine         *machine.Machine
	claimTTL        time.Duration
	refreshInterval time.Duration
	watches         map[string]job.JobWatch
	machines        map[string]machine.Machine
}

func NewJobWatcher(reg *registry.Registry, scheduler *Scheduler, m *machine.Machine) *JobWatcher {
	claimTTL, _ := time.ParseDuration(DefaultJobWatchClaimTTL)
	refreshInterval, _ := time.ParseDuration(DefaultRefreshInterval)
	return &JobWatcher{reg, scheduler, m, claimTTL, refreshInterval, make(map[string]job.JobWatch, 0), make(map[string]machine.Machine, 0)}
}

func (self *JobWatcher) StartHeartbeatThread() {
	heartbeat := func() {
		for _, watch := range self.watches {
			self.registry.ClaimJobWatch(&watch, self.machine, self.claimTTL)
		}
	}

	loop := func() {
		for true {
			heartbeat()
			time.Sleep(self.claimTTL / 2)
		}
	}

	go loop()
}

func (self *JobWatcher) StartRefreshThread() {
	refresh := func() {
		machines := make(map[string]machine.Machine, 0)
		for _, m := range self.registry.GetActiveMachines() {
			machines[m.BootId] = m
		}
		self.machines = machines
	}

	loop := func() {
		for true {
			refresh()
			time.Sleep(self.refreshInterval)
		}
	}

	go loop()
}

func (self *JobWatcher) AddJobWatch(watch *job.JobWatch) bool {
	if !self.registry.ClaimJobWatch(watch, self.machine, self.claimTTL) {
		return false
	}

	self.watches[watch.Payload.Name] = *watch
	sched := NewSchedule()

	if watch.Count == -1 {
		for _, m := range self.machines {
			name := fmt.Sprintf("%s.%s", m.BootId, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("EventJobWatchCreated(%s): adding to schedule job=%s machine=%s", watch.Payload.Name, name, m.BootId)
			sched.Add(*j, m)
		}
	} else {
		for i := 0; i < watch.Count; i++ {
			m := pickRandomMachine(self.machines)
			name := fmt.Sprintf("%s.%s", m.BootId, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("EventJobWatchCreated(%s): adding to schedule job=%s machine=%s", watch.Payload.Name, name, m.BootId)
			sched.Add(*j, *m)
		}
	}

	if len(sched) > 0 {
		log.Printf("EventJobWatchCreated(%s): submitting schedule", watch.Payload.Name)
		self.submitSchedule(sched)
	} else {
		log.Printf("EventJobWatchCreated(%s): no schedule changes made", watch.Payload.Name)
	}

	return true
}

func (self *JobWatcher) RemoveJobWatch(watch *job.JobWatch) bool {
	if _, ok := self.watches[watch.Payload.Name]; !ok {
		return false
	}

	delete(self.watches, watch.Payload.Name)
	return true
}

func (self *JobWatcher) submitSchedule(schedule Schedule) {
	for j, m := range schedule {
		self.registry.ScheduleMachineJob(&j, m)
	}
}

func (self *JobWatcher) TrackMachine(m *machine.Machine) {
	self.machines[m.BootId] = *m

	sched := NewSchedule()
	for _, watch := range self.watches {
		if watch.Count == -1 {
			name := fmt.Sprintf("%s.%s", m.BootId, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("Adding to schedule job=%s machine=%s", name, m.BootId)
			sched.Add(*j, *m)
		}
	}

	if len(sched) > 0 {
		log.Printf("Submitting schedule")
		self.submitSchedule(sched)
	} else {
		log.Printf("No schedule changes made")
	}
}

func (self *JobWatcher) DropMachine(m *machine.Machine) {
	if _, ok := self.machines[m.BootId]; ok {
		delete(self.machines, m.BootId)
	}
}
<commit_msg>refactor(JobWatcher): track schedules for ease of Job deletion<commit_after>package engine

import (
	"fmt"
	"log"
	"time"

	"github.com/coreos/coreinit/job"
	"github.com/coreos/coreinit/machine"
	"github.com/coreos/coreinit/registry"
)

const (
	DefaultJobWatchClaimTTL = "10s"
	DefaultRefreshInterval  = "60s"
)

type JobWatcher struct {
	registry        *registry.Registry
	scheduler       *Scheduler
	machine         *machine.Machine
	claimTTL        time.Duration
	refreshInterval time.Duration
	watches         map[string]job.JobWatch
	schedules       map[string]Schedule
	machines        map[string]machine.Machine
}

func NewJobWatcher(reg *registry.Registry, scheduler *Scheduler, m *machine.Machine) *JobWatcher {
	claimTTL, _ := time.ParseDuration(DefaultJobWatchClaimTTL)
	refreshInterval, _ := time.ParseDuration(DefaultRefreshInterval)
	return &JobWatcher{reg, scheduler, m, claimTTL, refreshInterval, make(map[string]job.JobWatch, 0),make(map[string]Schedule, 0), make(map[string]machine.Machine, 0)}
}

func (self *JobWatcher) StartHeartbeatThread() {
	heartbeat := func() {
		for _, watch := range self.watches {
			self.registry.ClaimJobWatch(&watch, self.machine, self.claimTTL)
		}
	}

	loop := func() {
		for true {
			heartbeat()
			time.Sleep(self.claimTTL / 2)
		}
	}

	go loop()
}

func (self *JobWatcher) StartRefreshThread() {
	refresh := func() {
		machines := make(map[string]machine.Machine, 0)
		for _, m := range self.registry.GetActiveMachines() {
			machines[m.BootId] = m
		}
		self.machines = machines
	}

	loop := func() {
		for true {
			refresh()
			time.Sleep(self.refreshInterval)
		}
	}

	go loop()
}

func (self *JobWatcher) AddJobWatch(watch *job.JobWatch) bool {
	if !self.registry.ClaimJobWatch(watch, self.machine, self.claimTTL) {
		return false
	}

	self.watches[watch.Payload.Name] = *watch
	sched := NewSchedule()
	self.schedules[watch.Payload.Name] = sched

	if watch.Count == -1 {
		for _, m := range self.machines {
			name := fmt.Sprintf("%s.%s", m.BootId, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("EventJobWatchCreated(%s): adding to schedule job=%s machine=%s", watch.Payload.Name, name, m.BootId)
			sched.Add(*j, m)
		}
	} else {
		for i := 1; i <= watch.Count; i++ {
			m := pickRandomMachine(self.machines)
			name := fmt.Sprintf("%d.%s", i, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("EventJobWatchCreated(%s): adding to schedule job=%s machine=%s", watch.Payload.Name, name, m.BootId)
			sched.Add(*j, *m)
		}
	}

	if len(sched) > 0 {
		log.Printf("EventJobWatchCreated(%s): submitting schedule", watch.Payload.Name)
		self.submitSchedule(sched)
	} else {
		log.Printf("EventJobWatchCreated(%s): no schedule changes made", watch.Payload.Name)
	}

	return true
}

func (self *JobWatcher) RemoveJobWatch(watch *job.JobWatch) bool {
	if _, ok := self.watches[watch.Payload.Name]; !ok {
		return false
	}

	delete(self.watches, watch.Payload.Name)

	watchSchedule := self.schedules[watch.Payload.Name]
	delete(self.schedules, watch.Payload.Name)

	for job, mach := range watchSchedule {
		self.registry.RemoveMachineJob(&job, mach)
	}

	return true
}

func (self *JobWatcher) submitSchedule(schedule Schedule) {
	for j, m := range schedule {
		self.registry.ScheduleMachineJob(&j, m)
	}
}

func (self *JobWatcher) TrackMachine(m *machine.Machine) {
	self.machines[m.BootId] = *m

	partial := NewSchedule()
	for _, watch := range self.watches {
		if watch.Count == -1 {
			name := fmt.Sprintf("%s.%s", m.BootId, watch.Payload.Name)
			j, _ := job.NewJob(name, nil, watch.Payload)
			log.Printf("Adding to schedule job=%s machine=%s", name, m.BootId)
			partial.Add(*j, *m)

			sched := self.schedules[watch.Payload.Name]
			sched.Add(*j, *m)
		}
	}

	if len(partial) > 0 {
		log.Printf("Submitting schedule")
		self.submitSchedule(partial)
	} else {
		log.Printf("No schedule changes made")
	}
}

func (self *JobWatcher) DropMachine(m *machine.Machine) {
	if _, ok := self.machines[m.BootId]; ok {
		delete(self.machines, m.BootId)
	}
}
<|endoftext|>


<commit_before>package operatingsystem

import (
  "io/ioutil"
  "os"
  "strings"
)

func Load() string {

    // redhat/centos/scientific
    // if /etc/redhat-release exists then we're dealing with redhat, centos, or scientific
	if _, err := os.Stat("/etc/redhat-release"); err == nil {

	    results, err := ioutil.ReadFile("/etc/redhat-release")
	        if err != nil {
	            panic(err)
	        }
	    return(string(results))

	}

    // ubuntu
    // if /etc/lsb-release exists then we're dealing with an ubuntu host
    if _, err := os.Stat("/etc/lsb-release"); err == nil {

        // Read in /etc/lsb-release
        results, err := ioutil.ReadFile("/etc/lsb-release")
        if err != nil {
            panic(err)
        }

        // scan through results w/ strings.Split
        for _, line :=  range strings.Split(string(results), "\n") {
            s := strings.Split(line, "=")
            if s[0] == "DISTRIB_DESCRIPTION" {
              return(strings.Replace(s[1], `"`, "", -1))
            }
        }
    }

	return(string("error: cannot detect operating system"))
}

<commit_msg>fixing operating system to return just the OS name for redhat family OSes<commit_after>package operatingsystem

import (
  "github.com/poblahblahblah/gofigure/lib/factfuncts"
  "io/ioutil"
  "os"
  "regexp"
  "strings"
)

func Load() string {

  // try to detect which OS we are
  // if /etc/redhat-release exists then we're dealing with redhat, centos, or scientific
  if _, err := os.Stat("/etc/redhat-release"); err == nil {

    centos_regexp, err     := regexp.Compile(`(?i)centos`)
    fedora_regexp, err     := regexp.Compile(`(?i)fedora`)
    redhat_regexp, err     := regexp.Compile(`(?i)redhat`)
    scientific_regexp, err := regexp.Compile(`(?i)scientific`)

    results, err := ioutil.ReadFile("/etc/redhat-release")
    if err != nil {
      panic(err)
    }

    os_long_string := factfuncts.Chomp(string(results))

    if centos_regexp.MatchString(os_long_string) == true {
      return string("CentOS")
    } else if redhat_regexp.MatchString(os_long_string) == true {
      return string("RedHat")
    } else if scientific_regexp.MatchString(os_long_string) == true {
      return string("Scientific Linux")
    }
  }

  // if /etc/debian_release exists, then we're using debian or a derivative
  if _, err := os.Stat("/etc/debian_release"); err == nil {

    // if /etc/lsb-release
    if _, err := os.Stat("/etc/lsb-release"); err == nil {

      // Read in /etc/lsb-release
      results, err := ioutil.ReadFile("/etc/lsb-release")
      if err != nil {
        panic(err)
      }

      // scan through results w/ strings.Split
      for _, line :=  range strings.Split(string(results), "\n") {
        s := strings.Split(line, "=")
        if s[0] == "DISTRIB_DESCRIPTION" {
          return(strings.Replace(s[1], `"`, "", -1))
        }
      }
    }
  }

  // FIXME:
  // Add Solaris, the BSDs, etc.

  return(string("error: cannot detect operating system"))
}

<|endoftext|>


<commit_before>package llog

import (
	"bytes"
	"fmt"
	"testing"
)

func TestOnlyErrorLogging(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, ERROR)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf")
}

func TestWarningOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, WARNING)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf")
}

func TestInfoOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, INFO)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf", "Info", "Infof")
}

func TestDebugOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, DEBUG)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf", "Info", "Infof", "Debug", "Debugf")
}

func logAndAssertNoError(t *testing.T, fn func(...interface{}) error, msg ...interface{}) {
	err := fn(msg...)
	if err != nil {
		t.Fatal(err)
	}
}

func logfAndAssertNoError(t *testing.T, fn func(string, ...interface{}) error, fmt string, msg ...interface{}) {
	err := fn(fmt, msg...)
	if err != nil {
		t.Fatal(err)
	}
}

func verifyLogEntries(t *testing.T, b *bytes.Buffer, entries ...string) {
	eb := new(bytes.Buffer)

	for _, e := range entries {
		eb.WriteString(fmt.Sprintln(e))
	}

	if eb.String() != b.String() {
		t.Errorf("Expected entries log does not match\nExpected: %s\nGot: %s\n", eb.String(), b.String())
	}
}
<commit_msg>Add test for SetLevel<commit_after>package llog

import (
	"bytes"
	"fmt"
	"testing"
)

func TestOnlyErrorLogging(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, ERROR)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf")
}

func TestWarningOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, WARNING)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf")
}

func TestInfoOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, INFO)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf", "Info", "Infof")
}

func TestDebugOrHigher(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, DEBUG)

	logAndAssertNoError(t, l.Error, "Error")
	logfAndAssertNoError(t, l.Errorf, "%s", "Errorf")

	logAndAssertNoError(t, l.Warning, "Warning")
	logfAndAssertNoError(t, l.Warningf, "%s", "Warningf")

	logAndAssertNoError(t, l.Info, "Info")
	logfAndAssertNoError(t, l.Infof, "%s", "Infof")

	logAndAssertNoError(t, l.Debug, "Debug")
	logfAndAssertNoError(t, l.Debugf, "%s", "Debugf")

	verifyLogEntries(t, b, "Error", "Errorf", "Warning", "Warningf", "Info", "Infof", "Debug", "Debugf")
}

func TestSetLevel(t *testing.T) {
	b := new(bytes.Buffer)
	l := New(b, INFO)

	logAndAssertNoError(t, l.Debug, "Error")
	l.SetLevel(DEBUG)
	logAndAssertNoError(t, l.Debug, "Error")

	verifyLogEntries(t, b, "Error")
}

func logAndAssertNoError(t *testing.T, fn func(...interface{}) error, msg ...interface{}) {
	err := fn(msg...)
	if err != nil {
		t.Fatal(err)
	}
}

func logfAndAssertNoError(t *testing.T, fn func(string, ...interface{}) error, fmt string, msg ...interface{}) {
	err := fn(fmt, msg...)
	if err != nil {
		t.Fatal(err)
	}
}

func verifyLogEntries(t *testing.T, b *bytes.Buffer, entries ...string) {
	eb := new(bytes.Buffer)

	for _, e := range entries {
		eb.WriteString(fmt.Sprintln(e))
	}

	if eb.String() != b.String() {
		t.Errorf("Expected entries log does not match\nExpected: %s\nGot: %s\n", eb.String(), b.String())
	}
}
<|endoftext|>


<commit_before>// Code generated by protoc-gen-yarpc-go. DO NOT EDIT.
// source: echo.proto

package echo

import (
	"context"
	"io/ioutil"
	"reflect"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"go.uber.org/fx"
	"go.uber.org/yarpc"
	"go.uber.org/yarpc/api/transport"
	"go.uber.org/yarpc/encoding/protobuf"
	"go.uber.org/yarpc/encoding/protobuf/reflection"
)

var _ = ioutil.NopCloser

// EchoYARPCClient is the YARPC client-side interface for the Echo service.
type EchoYARPCClient interface {
	Echo(context.Context, *Request, ...yarpc.CallOption) (*Response, error)
}

func newEchoYARPCClient(clientConfig transport.ClientConfig, anyResolver jsonpb.AnyResolver, options ...protobuf.ClientOption) EchoYARPCClient {
	return &_EchoYARPCCaller{protobuf.NewStreamClient(
		protobuf.ClientParams{
			ServiceName:  "echo.Echo",
			ClientConfig: clientConfig,
			AnyResolver:  anyResolver,
			Options:      options,
		},
	)}
}

// NewEchoYARPCClient builds a new YARPC client for the Echo service.
func NewEchoYARPCClient(clientConfig transport.ClientConfig, options ...protobuf.ClientOption) EchoYARPCClient {
	return newEchoYARPCClient(clientConfig, nil, options...)
}

// EchoYARPCServer is the YARPC server-side interface for the Echo service.
type EchoYARPCServer interface {
	Echo(context.Context, *Request) (*Response, error)
}

type buildEchoYARPCProceduresParams struct {
	Server      EchoYARPCServer
	AnyResolver jsonpb.AnyResolver
}

func buildEchoYARPCProcedures(params buildEchoYARPCProceduresParams) []transport.Procedure {
	handler := &_EchoYARPCHandler{params.Server}
	return protobuf.BuildProcedures(
		protobuf.BuildProceduresParams{
			ServiceName: "echo.Echo",
			UnaryHandlerParams: []protobuf.BuildProceduresUnaryHandlerParams{
				{
					MethodName: "Echo",
					Handler: protobuf.NewUnaryHandler(
						protobuf.UnaryHandlerParams{
							Handle:      handler.Echo,
							NewRequest:  newEchoServiceEchoYARPCRequest,
							AnyResolver: params.AnyResolver,
						},
					),
				},
			},
			OnewayHandlerParams: []protobuf.BuildProceduresOnewayHandlerParams{},
			StreamHandlerParams: []protobuf.BuildProceduresStreamHandlerParams{},
		},
	)
}

// BuildEchoYARPCProcedures prepares an implementation of the Echo service for YARPC registration.
func BuildEchoYARPCProcedures(server EchoYARPCServer) []transport.Procedure {
	return buildEchoYARPCProcedures(buildEchoYARPCProceduresParams{Server: server})
}

// FxEchoYARPCClientParams defines the input
// for NewFxEchoYARPCClient. It provides the
// paramaters to get a EchoYARPCClient in an
// Fx application.
type FxEchoYARPCClientParams struct {
	fx.In

	Provider    yarpc.ClientConfig
	AnyResolver jsonpb.AnyResolver `name:"yarpcfx" optional:"true"`
}

// FxEchoYARPCClientResult defines the output
// of NewFxEchoYARPCClient. It provides a
// EchoYARPCClient to an Fx application.
type FxEchoYARPCClientResult struct {
	fx.Out

	Client EchoYARPCClient

	// We are using an fx.Out struct here instead of just returning a client
	// so that we can add more values or add named versions of the client in
	// the future without breaking any existing code.
}

// NewFxEchoYARPCClient provides a EchoYARPCClient
// to an Fx application using the given name for routing.
//
//  fx.Provide(
//    echo.NewFxEchoYARPCClient("service-name"),
//    ...
//  )
func NewFxEchoYARPCClient(name string, options ...protobuf.ClientOption) interface{} {
	return func(params FxEchoYARPCClientParams) FxEchoYARPCClientResult {
		return FxEchoYARPCClientResult{
			Client: newEchoYARPCClient(params.Provider.ClientConfig(name), params.AnyResolver, options...),
		}
	}
}

// FxEchoYARPCProceduresParams defines the input
// for NewFxEchoYARPCProcedures. It provides the
// paramaters to get EchoYARPCServer procedures in an
// Fx application.
type FxEchoYARPCProceduresParams struct {
	fx.In

	Server      EchoYARPCServer
	AnyResolver jsonpb.AnyResolver `name:"yarpcfx" optional:"true"`
}

// FxEchoYARPCProceduresResult defines the output
// of NewFxEchoYARPCProcedures. It provides
// EchoYARPCServer procedures to an Fx application.
//
// The procedures are provided to the "yarpcfx" value group.
// Dig 1.2 or newer must be used for this feature to work.
type FxEchoYARPCProceduresResult struct {
	fx.Out

	Procedures     []transport.Procedure `group:"yarpcfx"`
	ReflectionMeta reflection.ServerMeta `group:"yarpcfx"`
}

// NewFxEchoYARPCProcedures provides EchoYARPCServer procedures to an Fx application.
// It expects a EchoYARPCServer to be present in the container.
//
//  fx.Provide(
//    echo.NewFxEchoYARPCProcedures(),
//    ...
//  )
func NewFxEchoYARPCProcedures() interface{} {
	return func(params FxEchoYARPCProceduresParams) FxEchoYARPCProceduresResult {
		return FxEchoYARPCProceduresResult{
			Procedures: buildEchoYARPCProcedures(buildEchoYARPCProceduresParams{
				Server:      params.Server,
				AnyResolver: params.AnyResolver,
			}),
			ReflectionMeta: reflection.ServerMeta{
				ServiceName:     "echo.Echo",
				FileDescriptors: yarpcFileDescriptorClosure08134aea513e0001,
			},
		}
	}
}

type _EchoYARPCCaller struct {
	streamClient protobuf.StreamClient
}

func (c *_EchoYARPCCaller) Echo(ctx context.Context, request *Request, options ...yarpc.CallOption) (*Response, error) {
	responseMessage, err := c.streamClient.Call(ctx, "Echo", request, newEchoServiceEchoYARPCResponse, options...)
	if responseMessage == nil {
		return nil, err
	}
	response, ok := responseMessage.(*Response)
	if !ok {
		return nil, protobuf.CastError(emptyEchoServiceEchoYARPCResponse, responseMessage)
	}
	return response, err
}

type _EchoYARPCHandler struct {
	server EchoYARPCServer
}

func (h *_EchoYARPCHandler) Echo(ctx context.Context, requestMessage proto.Message) (proto.Message, error) {
	var request *Request
	var ok bool
	if requestMessage != nil {
		request, ok = requestMessage.(*Request)
		if !ok {
			return nil, protobuf.CastError(emptyEchoServiceEchoYARPCRequest, requestMessage)
		}
	}
	response, err := h.server.Echo(ctx, request)
	if response == nil {
		return nil, err
	}
	return response, err
}

func newEchoServiceEchoYARPCRequest() proto.Message {
	return &Request{}
}

func newEchoServiceEchoYARPCResponse() proto.Message {
	return &Response{}
}

var (
	emptyEchoServiceEchoYARPCRequest  = &Request{}
	emptyEchoServiceEchoYARPCResponse = &Response{}
)

var yarpcFileDescriptorClosure08134aea513e0001 = [][]byte{
	// echo.proto
	[]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xe2, 0xe2, 0x4a, 0x4d, 0xce, 0xc8,
		0xd7, 0x2b, 0x28, 0xca, 0x2f, 0xc9, 0x17, 0x62, 0x01, 0xb1, 0x95, 0x94, 0xb9, 0xd8, 0x83, 0x52,
		0x0b, 0x4b, 0x53, 0x8b, 0x4b, 0x84, 0x24, 0xb8, 0xd8, 0x73, 0x53, 0x8b, 0x8b, 0x13, 0xd3, 0x53,
		0x25, 0x18, 0x15, 0x18, 0x35, 0x38, 0x83, 0x60, 0x5c, 0x25, 0x15, 0x2e, 0x8e, 0xa0, 0xd4, 0xe2,
		0x82, 0xfc, 0xbc, 0xe2, 0x54, 0xdc, 0xaa, 0x8c, 0x74, 0xb9, 0x58, 0x5c, 0x93, 0x33, 0xf2, 0x85,
		0x54, 0xa1, 0x34, 0xaf, 0x1e, 0xd8, 0x36, 0xa8, 0xf1, 0x52, 0x7c, 0x30, 0x2e, 0xc4, 0xa0, 0x24,
		0x36, 0xb0, 0x33, 0x8c, 0x01, 0x01, 0x00, 0x00, 0xff, 0xff, 0xe1, 0x1c, 0xf0, 0xfd, 0x94, 0x00,
		0x00, 0x00,
	},
}

func init() {
	yarpc.RegisterClientBuilder(
		func(clientConfig transport.ClientConfig, structField reflect.StructField) EchoYARPCClient {
			return NewEchoYARPCClient(clientConfig, protobuf.ClientBuilderOptions(clientConfig, structField)...)
		},
	)
}
<commit_msg>Rebasing master and standardizing goimports<commit_after>// Code generated by protoc-gen-yarpc-go. DO NOT EDIT.
// source: clients/echo/echo.proto

package echo

import (
	"context"
	"io/ioutil"
	"reflect"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"go.uber.org/fx"
	"go.uber.org/yarpc"
	"go.uber.org/yarpc/api/transport"
	"go.uber.org/yarpc/encoding/protobuf"
	"go.uber.org/yarpc/encoding/protobuf/reflection"
)

var _ = ioutil.NopCloser

// EchoYARPCClient is the YARPC client-side interface for the Echo service.
type EchoYARPCClient interface {
	Echo(context.Context, *Request, ...yarpc.CallOption) (*Response, error)
}

func newEchoYARPCClient(clientConfig transport.ClientConfig, anyResolver jsonpb.AnyResolver, options ...protobuf.ClientOption) EchoYARPCClient {
	return &_EchoYARPCCaller{protobuf.NewStreamClient(
		protobuf.ClientParams{
			ServiceName:  "echo.Echo",
			ClientConfig: clientConfig,
			AnyResolver:  anyResolver,
			Options:      options,
		},
	)}
}

// NewEchoYARPCClient builds a new YARPC client for the Echo service.
func NewEchoYARPCClient(clientConfig transport.ClientConfig, options ...protobuf.ClientOption) EchoYARPCClient {
	return newEchoYARPCClient(clientConfig, nil, options...)
}

// EchoYARPCServer is the YARPC server-side interface for the Echo service.
type EchoYARPCServer interface {
	Echo(context.Context, *Request) (*Response, error)
}

type buildEchoYARPCProceduresParams struct {
	Server      EchoYARPCServer
	AnyResolver jsonpb.AnyResolver
}

func buildEchoYARPCProcedures(params buildEchoYARPCProceduresParams) []transport.Procedure {
	handler := &_EchoYARPCHandler{params.Server}
	return protobuf.BuildProcedures(
		protobuf.BuildProceduresParams{
			ServiceName: "echo.Echo",
			UnaryHandlerParams: []protobuf.BuildProceduresUnaryHandlerParams{
				{
					MethodName: "Echo",
					Handler: protobuf.NewUnaryHandler(
						protobuf.UnaryHandlerParams{
							Handle:      handler.Echo,
							NewRequest:  newEchoServiceEchoYARPCRequest,
							AnyResolver: params.AnyResolver,
						},
					),
				},
			},
			OnewayHandlerParams: []protobuf.BuildProceduresOnewayHandlerParams{},
			StreamHandlerParams: []protobuf.BuildProceduresStreamHandlerParams{},
		},
	)
}

// BuildEchoYARPCProcedures prepares an implementation of the Echo service for YARPC registration.
func BuildEchoYARPCProcedures(server EchoYARPCServer) []transport.Procedure {
	return buildEchoYARPCProcedures(buildEchoYARPCProceduresParams{Server: server})
}

// FxEchoYARPCClientParams defines the input
// for NewFxEchoYARPCClient. It provides the
// paramaters to get a EchoYARPCClient in an
// Fx application.
type FxEchoYARPCClientParams struct {
	fx.In

	Provider    yarpc.ClientConfig
	AnyResolver jsonpb.AnyResolver `name:"yarpcfx" optional:"true"`
}

// FxEchoYARPCClientResult defines the output
// of NewFxEchoYARPCClient. It provides a
// EchoYARPCClient to an Fx application.
type FxEchoYARPCClientResult struct {
	fx.Out

	Client EchoYARPCClient

	// We are using an fx.Out struct here instead of just returning a client
	// so that we can add more values or add named versions of the client in
	// the future without breaking any existing code.
}

// NewFxEchoYARPCClient provides a EchoYARPCClient
// to an Fx application using the given name for routing.
//
//  fx.Provide(
//    echo.NewFxEchoYARPCClient("service-name"),
//    ...
//  )
func NewFxEchoYARPCClient(name string, options ...protobuf.ClientOption) interface{} {
	return func(params FxEchoYARPCClientParams) FxEchoYARPCClientResult {
		return FxEchoYARPCClientResult{
			Client: newEchoYARPCClient(params.Provider.ClientConfig(name), params.AnyResolver, options...),
		}
	}
}

// FxEchoYARPCProceduresParams defines the input
// for NewFxEchoYARPCProcedures. It provides the
// paramaters to get EchoYARPCServer procedures in an
// Fx application.
type FxEchoYARPCProceduresParams struct {
	fx.In

	Server      EchoYARPCServer
	AnyResolver jsonpb.AnyResolver `name:"yarpcfx" optional:"true"`
}

// FxEchoYARPCProceduresResult defines the output
// of NewFxEchoYARPCProcedures. It provides
// EchoYARPCServer procedures to an Fx application.
//
// The procedures are provided to the "yarpcfx" value group.
// Dig 1.2 or newer must be used for this feature to work.
type FxEchoYARPCProceduresResult struct {
	fx.Out

	Procedures     []transport.Procedure `group:"yarpcfx"`
	ReflectionMeta reflection.ServerMeta `group:"yarpcfx"`
}

// NewFxEchoYARPCProcedures provides EchoYARPCServer procedures to an Fx application.
// It expects a EchoYARPCServer to be present in the container.
//
//  fx.Provide(
//    echo.NewFxEchoYARPCProcedures(),
//    ...
//  )
func NewFxEchoYARPCProcedures() interface{} {
	return func(params FxEchoYARPCProceduresParams) FxEchoYARPCProceduresResult {
		return FxEchoYARPCProceduresResult{
			Procedures: buildEchoYARPCProcedures(buildEchoYARPCProceduresParams{
				Server:      params.Server,
				AnyResolver: params.AnyResolver,
			}),
			ReflectionMeta: reflection.ServerMeta{
				ServiceName:     "echo.Echo",
				FileDescriptors: yarpcFileDescriptorClosure6caa5dd77ae15c91,
			},
		}
	}
}

type _EchoYARPCCaller struct {
	streamClient protobuf.StreamClient
}

func (c *_EchoYARPCCaller) Echo(ctx context.Context, request *Request, options ...yarpc.CallOption) (*Response, error) {
	responseMessage, err := c.streamClient.Call(ctx, "Echo", request, newEchoServiceEchoYARPCResponse, options...)
	if responseMessage == nil {
		return nil, err
	}
	response, ok := responseMessage.(*Response)
	if !ok {
		return nil, protobuf.CastError(emptyEchoServiceEchoYARPCResponse, responseMessage)
	}
	return response, err
}

type _EchoYARPCHandler struct {
	server EchoYARPCServer
}

func (h *_EchoYARPCHandler) Echo(ctx context.Context, requestMessage proto.Message) (proto.Message, error) {
	var request *Request
	var ok bool
	if requestMessage != nil {
		request, ok = requestMessage.(*Request)
		if !ok {
			return nil, protobuf.CastError(emptyEchoServiceEchoYARPCRequest, requestMessage)
		}
	}
	response, err := h.server.Echo(ctx, request)
	if response == nil {
		return nil, err
	}
	return response, err
}

func newEchoServiceEchoYARPCRequest() proto.Message {
	return &Request{}
}

func newEchoServiceEchoYARPCResponse() proto.Message {
	return &Response{}
}

var (
	emptyEchoServiceEchoYARPCRequest  = &Request{}
	emptyEchoServiceEchoYARPCResponse = &Response{}
)

var yarpcFileDescriptorClosure6caa5dd77ae15c91 = [][]byte{
	// clients/echo/echo.proto
	[]byte{
		0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xe2, 0x12, 0x4f, 0xce, 0xc9, 0x4c,
		0xcd, 0x2b, 0x29, 0xd6, 0x4f, 0x4d, 0xce, 0xc8, 0x07, 0x13, 0x7a, 0x05, 0x45, 0xf9, 0x25, 0xf9,
		0x42, 0x2c, 0x20, 0xb6, 0x92, 0x32, 0x17, 0x7b, 0x50, 0x6a, 0x61, 0x69, 0x6a, 0x71, 0x89, 0x90,
		0x04, 0x17, 0x7b, 0x6e, 0x6a, 0x71, 0x71, 0x62, 0x7a, 0xaa, 0x04, 0xa3, 0x02, 0xa3, 0x06, 0x67,
		0x10, 0x8c, 0xab, 0xa4, 0xc2, 0xc5, 0x11, 0x94, 0x5a, 0x5c, 0x90, 0x9f, 0x57, 0x9c, 0x8a, 0x5b,
		0x95, 0x91, 0x2e, 0x17, 0x8b, 0x6b, 0x72, 0x46, 0xbe, 0x90, 0x2a, 0x94, 0xe6, 0xd5, 0x03, 0xdb,
		0x06, 0x35, 0x5e, 0x8a, 0x0f, 0xc6, 0x85, 0x18, 0x94, 0xc4, 0x06, 0x76, 0x86, 0x31, 0x20, 0x00,
		0x00, 0xff, 0xff, 0xe9, 0xcf, 0xa3, 0x1b, 0xa1, 0x00, 0x00, 0x00,
	},
}

func init() {
	yarpc.RegisterClientBuilder(
		func(clientConfig transport.ClientConfig, structField reflect.StructField) EchoYARPCClient {
			return NewEchoYARPCClient(clientConfig, protobuf.ClientBuilderOptions(clientConfig, structField)...)
		},
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package e2e

import (
	"fmt"
	"time"

	"k8s.io/frakti/test/e2e/framework"
	internalapi "k8s.io/kubernetes/pkg/kubelet/api"
	kubeapi "k8s.io/kubernetes/pkg/kubelet/api/v1alpha1/runtime"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = framework.KubeDescribe("Test streaming in container", func() {
	f := framework.NewDefaultFramework("test")

	var (
		runtimeClient internalapi.RuntimeService
		imageClient   internalapi.ImageManagerService
	)

	BeforeEach(func() {
		runtimeClient = f.Client.FraktiRuntimeService
		imageClient = f.Client.FraktiImageService
	})

	It("test exec a command in container synchronously and successfully", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podId string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("exec command in container synchronously")
		magicWords := "blablabla"
		stdout, stderr, err := runtimeClient.ExecSync(cID, []string{"echo", magicWords}, 0)
		framework.ExpectNoError(err, "Failed to exec cmd in container: %v", err)
		framework.Logf("stdout: %q, stderr: %q", string(stdout), string(stderr))
		Expect(len(stderr)).To(Equal(0), "stderr should not have content")
		Expect(string(stdout)).To(BeIdenticalTo(magicWords+"\n"), "stdout should be same as defined")
	})

	It("test exec a command in container synchronously and failed", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podID string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("exec command in container synchronously")
		magicCmd := "blablabla"
		stdout, stderr, err := runtimeClient.ExecSync(cID, []string{magicCmd}, 0)
		Expect(err).NotTo(Equal(nil), "Exec non-exist cmd should failed")
		framework.Logf("stdout: %q, stderr: %q", string(stdout), string(stderr))
		Expect(len(stderr)).NotTo(Equal(0), "stderr should have content")
	})

	It("test get a exec url", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podID string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("prepare exec command url in container")
		magicCmd := []string{"blablabla"}
		execReq := &kubeapi.ExecRequest{
			ContainerId: &cID,
			Cmd:         magicCmd,
		}
		resp, err := runtimeClient.Exec(execReq)
		framework.ExpectNoError(err, "Failed to get exec url in container: %v", err)
		framework.Logf("ExecUrl: %q", resp.GetUrl())
		expectedUrl := buildExpectedExecAttachURL("exec", cID, magicCmd, false, false)
		Expect(resp.GetUrl()).To(Equal(expectedUrl), "exec url should equal to expected")
	})
})

func startLongRunningContainer(rc internalapi.RuntimeService, ic internalapi.ImageManagerService) (podId, containerId string) {
	podName := "simple-sandbox-" + framework.NewUUID()
	By("create a podSandbox")
	podConfig := &kubeapi.PodSandboxConfig{
		Metadata: &kubeapi.PodSandboxMetadata{
			Name: &podName,
		},
	}
	podId, err := rc.RunPodSandbox(podConfig)
	framework.ExpectNoError(err, "Failed to create podsandbox: %v", err)

	By("pull necessary image")
	imageSpec := &kubeapi.ImageSpec{
		Image: &latestTestImageRef,
	}
	err = ic.PullImage(imageSpec, nil)
	framework.ExpectNoError(err, "Failed to pull image: %v", err)

	By("create container in pod")
	containerName := "simple-container-" + framework.NewUUID()
	containerConfig := &kubeapi.ContainerConfig{
		Metadata: &kubeapi.ContainerMetadata{
			Name: &containerName,
		},
		Image:   imageSpec,
		Command: []string{"sh", "-c", "top"},
	}
	containerId, err = rc.CreateContainer(podId, containerConfig, podConfig)
	framework.ExpectNoError(err, "Failed to create container: %v", err)

	By("start container")
	err = rc.StartContainer(containerId)
	framework.ExpectNoError(err, "Failed to start container: %v", err)

	// sleep 2s to make sure container start is ready, workaround for random failed in travis
	// TODO: remove this
	time.Sleep(2 * time.Second)

	return podId, containerId
}

func buildExpectedExecAttachURL(method, containerID string, cmd []string, stdin, tty bool) string {
	query := ""
	// build cmd
	for i, c := range cmd {
		if i == 0 {
			query += "?"
		} else {
			query += "&"
		}
		query += fmt.Sprintf("command=%s", c)
	}
	if len(cmd) == 0 {
		query += "?"
	} else {
		query += "&"
	}
	// build io setting
	if !tty {
		query += "error=1&"
	}
	if stdin {
		query += "input=1&"
	}
	// stdout is set default
	query += "output=1&"
	if tty {
		query += "tty=1&"
	}
	// remove traval '&'
	query = query[:len(query)-1]

	// TODO: http schema and address should generate from frakti config
	return fmt.Sprintf("http://0.0.0.0:22521/%s/%s%s", method, containerID, query)
}
<commit_msg>add test for get attach url api<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package e2e

import (
	"fmt"
	"time"

	"k8s.io/frakti/test/e2e/framework"
	internalapi "k8s.io/kubernetes/pkg/kubelet/api"
	kubeapi "k8s.io/kubernetes/pkg/kubelet/api/v1alpha1/runtime"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = framework.KubeDescribe("Test streaming in container", func() {
	f := framework.NewDefaultFramework("test")

	var (
		runtimeClient internalapi.RuntimeService
		imageClient   internalapi.ImageManagerService
	)

	BeforeEach(func() {
		runtimeClient = f.Client.FraktiRuntimeService
		imageClient = f.Client.FraktiImageService
	})

	It("test exec a command in container synchronously and successfully", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podId string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("exec command in container synchronously")
		magicWords := "blablabla"
		stdout, stderr, err := runtimeClient.ExecSync(cID, []string{"echo", magicWords}, 0)
		framework.ExpectNoError(err, "Failed to exec cmd in container: %v", err)
		framework.Logf("stdout: %q, stderr: %q", string(stdout), string(stderr))
		Expect(len(stderr)).To(Equal(0), "stderr should not have content")
		Expect(string(stdout)).To(BeIdenticalTo(magicWords+"\n"), "stdout should be same as defined")
	})

	It("test exec a command in container synchronously and failed", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podID string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("exec command in container synchronously")
		magicCmd := "blablabla"
		stdout, stderr, err := runtimeClient.ExecSync(cID, []string{magicCmd}, 0)
		Expect(err).NotTo(Equal(nil), "Exec non-exist cmd should failed")
		framework.Logf("stdout: %q, stderr: %q", string(stdout), string(stderr))
		Expect(len(stderr)).NotTo(Equal(0), "stderr should have content")
	})

	It("test get a exec url", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podID string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("prepare exec command url in container")
		magicCmd := []string{"blablabla"}
		execReq := &kubeapi.ExecRequest{
			ContainerId: &cID,
			Cmd:         magicCmd,
		}
		resp, err := runtimeClient.Exec(execReq)
		framework.ExpectNoError(err, "Failed to get exec url in container: %v", err)
		framework.Logf("ExecUrl: %q", resp.GetUrl())
		expectedUrl := buildExpectedExecAttachURL("exec", cID, magicCmd, false, false)
		Expect(resp.GetUrl()).To(Equal(expectedUrl), "exec url should equal to expected")
	})

	It("test get a attach url", func() {
		podID, cID := startLongRunningContainer(runtimeClient, imageClient)
		defer func(podID string) {
			By("delete pod sandbox")
			runtimeClient.RemovePodSandbox(podID)
		}(podID)

		By("prepare attach command url in container")
		stdin := true
		attachReq := &kubeapi.AttachRequest{
			ContainerId: &cID,
			Stdin:       &stdin,
		}
		resp, err := runtimeClient.Attach(attachReq)
		framework.ExpectNoError(err, "Failed to get attach url in container: %v", err)
		framework.Logf("AttachUrl: %q", resp.GetUrl())
		expectedUrl := buildExpectedExecAttachURL("attach", cID, nil, stdin, false)
		Expect(resp.GetUrl()).To(Equal(expectedUrl), "attach url should equal to expected")
	})
})

func startLongRunningContainer(rc internalapi.RuntimeService, ic internalapi.ImageManagerService) (podId, containerId string) {
	podName := "simple-sandbox-" + framework.NewUUID()
	By("create a podSandbox")
	podConfig := &kubeapi.PodSandboxConfig{
		Metadata: &kubeapi.PodSandboxMetadata{
			Name: &podName,
		},
	}
	podId, err := rc.RunPodSandbox(podConfig)
	framework.ExpectNoError(err, "Failed to create podsandbox: %v", err)

	By("pull necessary image")
	imageSpec := &kubeapi.ImageSpec{
		Image: &latestTestImageRef,
	}
	err = ic.PullImage(imageSpec, nil)
	framework.ExpectNoError(err, "Failed to pull image: %v", err)

	By("create container in pod")
	containerName := "simple-container-" + framework.NewUUID()
	containerConfig := &kubeapi.ContainerConfig{
		Metadata: &kubeapi.ContainerMetadata{
			Name: &containerName,
		},
		Image:   imageSpec,
		Command: []string{"sh", "-c", "top"},
	}
	containerId, err = rc.CreateContainer(podId, containerConfig, podConfig)
	framework.ExpectNoError(err, "Failed to create container: %v", err)

	By("start container")
	err = rc.StartContainer(containerId)
	framework.ExpectNoError(err, "Failed to start container: %v", err)

	// sleep 2s to make sure container start is ready, workaround for random failed in travis
	// TODO: remove this
	time.Sleep(2 * time.Second)

	return podId, containerId
}

func buildExpectedExecAttachURL(method, containerID string, cmd []string, stdin, tty bool) string {
	query := ""
	// build cmd
	for i, c := range cmd {
		if i == 0 {
			query += "?"
		} else {
			query += "&"
		}
		query += fmt.Sprintf("command=%s", c)
	}
	if len(cmd) == 0 {
		query += "?"
	} else {
		query += "&"
	}
	// build io setting
	if !tty {
		query += "error=1&"
	}
	if stdin {
		query += "input=1&"
	}
	// stdout is set default
	query += "output=1&"
	if tty {
		query += "tty=1&"
	}
	// remove traval '&'
	query = query[:len(query)-1]

	// TODO: http schema and address should generate from frakti config
	return fmt.Sprintf("http://0.0.0.0:22521/%s/%s%s", method, containerID, query)
}
<|endoftext|>