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

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"sort"
	"sync"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
)

var fset = token.NewFileSet()

type Edit struct {
	Pos, End int
}

type editsByPos []Edit

func (e editsByPos) Len() int           { return len(e) }
func (e editsByPos) Less(i, j int) bool { return e[i].Pos < e[j].Pos }
func (e editsByPos) Swap(i, j int)      { e[i], e[j] = e[j], e[i] }

func apply(file string, edits []Edit) {
	if len(edits) == 0 {
		return
	}

	sort.Sort(editsByPos(edits))

	// Check for overlap.
	// TODO(mdempsky): Overlap can legally happen in bizarro expressions like
	// "(*[unsafe.Sizeof(int8(int8(0)))]byte)(*[1]byte)(nil)".
	for i := 1; i < len(edits); i++ {
		if edits[i-1].End > edits[i].Pos {
			log.Fatal("overlap")
		}
	}

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}

	n := edits[0].Pos
	for i := 1; i < len(edits); i++ {
		n += copy(buf[n:], buf[edits[i-1].End:edits[i].Pos])
	}
	n += copy(buf[n:], buf[edits[len(edits)-1].End:])
	buf = buf[:n]

	buf, err = format.Source(buf)
	if err != nil {
		log.Fatal(err)
	}

	err = ioutil.WriteFile(file, buf, 0)
	if err != nil {
		log.Fatal(err)
	}
}

var (
	flagAll   = flag.Bool("all", false, "type check all GOOS and GOARCH combinations")
	flagApply = flag.Bool("apply", false, "apply edits")
)

func main() {
	flag.Parse()

	var m map[string][]Edit
	if *flagAll {
		m = mergeEdits()
	} else {
		m = computeEdits(build.Default.GOOS, build.Default.GOARCH)
	}

	if *flagApply {
		for f, e := range m {
			apply(f, e)
		}
	} else {
		err := json.NewEncoder(os.Stdout).Encode(m)
		if err != nil {
			log.Fatal(err)
		}
		return
	}
}

var plats = [...]struct {
	goos, goarch string
}{
	{"linux", "386"},
	{"linux", "amd64"},
	{"linux", "arm"},
	{"linux", "arm64"},
	{"linux", "ppc64"},
	{"linux", "ppc64le"},
	{"nacl", "386"},
	{"nacl", "amd64p32"},
	{"nacl", "arm"},
	{"darwin", "386"},
	{"darwin", "amd64"},
	{"dragonfly", "amd64"},
	{"freebsd", "386"},
	{"freebsd", "amd64"},
	{"freebsd", "arm"},
	{"netbsd", "386"},
	{"netbsd", "amd64"},
	{"netbsd", "arm"},
	{"openbsd", "386"},
	{"openbsd", "amd64"},
	{"openbsd", "arm"},
	{"plan9", "386"},
	{"plan9", "amd64"},
	{"solaris", "amd64"},
	{"windows", "386"},
	{"windows", "amd64"},
}

func mergeEdits() map[string][]Edit {
	ch := make(chan map[string][]Edit)
	var wg sync.WaitGroup
	for _, plat := range plats {
		wg.Add(1)
		go func(goos, goarch string) {
			defer wg.Done()
			ch <- computeEdits(goos, goarch)
		}(plat.goos, plat.goarch)
	}
	go func() {
		wg.Wait()
		close(ch)
	}()

	m := make(map[string][]Edit)
	for m1 := range ch {
		for f, e := range m1 {
			if e0, ok := m[f]; ok {
				m[f] = intersect(e0, e)
			} else {
				m[f] = e
			}
		}
	}
	return m
}

func intersect(e1, e2 []Edit) []Edit {
	if len(e1) == 0 || len(e2) == 0 {
		return nil
	}

	set := make(map[Edit]bool, len(e1))
	for _, e := range e1 {
		set[e] = true
	}

	var res []Edit
	for _, e := range e2 {
		if set[e] {
			res = append(res, e)
		}
	}
	return res
}

func noImport(map[string]*types.Package, string) (*types.Package, error) {
	panic("go/loader said this wouldn't be called")
}

func computeEdits(os, arch string) map[string][]Edit {
	ctxt := build.Default
	ctxt.GOOS = os
	ctxt.GOARCH = arch
	ctxt.CgoEnabled = false

	var conf loader.Config
	conf.Fset = fset
	conf.Build = &ctxt
	conf.TypeChecker.Import = noImport
	for _, arg := range flag.Args() {
		conf.Import(arg)
	}
	prog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}

	type res struct {
		file  string
		edits []Edit
	}
	ch := make(chan res)
	var wg sync.WaitGroup
	for _, pkg := range prog.InitialPackages() {
		for _, file := range pkg.Files {
			wg.Add(1)
			go func(pkg *loader.PackageInfo, file *ast.File) {
				defer wg.Done()
				v := visitor{pkg: pkg, file: fset.File(file.Package)}
				ast.Walk(&v, file)
				sort.Sort(editsByPos(v.edits))
				ch <- res{v.file.Name(), v.edits}
			}(pkg, file)
		}
	}
	go func() {
		wg.Wait()
		close(ch)
	}()

	m := make(map[string][]Edit)
	for r := range ch {
		m[r.file] = r.edits
	}
	return m
}

type visitor struct {
	pkg   *loader.PackageInfo
	file  *token.File
	edits []Edit
	nodes []ast.Node
}

func (v *visitor) Visit(node ast.Node) ast.Visitor {
	if node != nil {
		v.nodes = append(v.nodes, node)
	} else {
		v.nodes = v.nodes[:len(v.nodes)-1]
	}
	if call, ok := node.(*ast.CallExpr); ok {
		v.unconvert(call)
	}
	return v
}

func (v *visitor) unconvert(call *ast.CallExpr) {
	// TODO(mdempsky): Handle useless multi-conversions.

	// Conversions have exactly one argument.
	if len(call.Args) != 1 || call.Ellipsis != token.NoPos {
		return
	}
	ft, ok := v.pkg.Types[call.Fun]
	if !ok {
		fmt.Println("Missing type for function")
		return
	}
	if !ft.IsType() {
		// Not a conversion.
		return
	}
	at, ok := v.pkg.Types[call.Args[0]]
	if !ok {
		fmt.Println("Missing type for argument")
	}
	if !types.Identical(ft.Type, at.Type) {
		// A real conversion.
		return
	}
	if at.Value != nil || hasUntypedValue(call.Args[0]) {
		// As a workaround for golang.org/issue/13061,
		// skip conversions that contain an untyped value.
		return
	}

	outer := v.nodes[len(v.nodes)-2]
	if keepParen(outer, call, call.Args[0]) {
		v.remove(call.Fun.Pos(), call.Lparen)
	} else {
		v.remove(call.Fun.Pos(), call.Lparen+1)
		v.remove(call.Rparen, call.Rparen+1)
	}
}

func (v *visitor) remove(pos, end token.Pos) {
	v.edits = append(v.edits, Edit{v.file.Offset(pos), v.file.Offset(end)})
}

func keepParen(a ast.Node, b, c ast.Expr) bool {
	// 1. Find the value in a that points to b.
	bp := findExprField(a, b)

	// 2. Try printing a with s/b/c/ and with s/b/(c)/.
	var buf1 bytes.Buffer
	*bp = c
	printer.Fprint(&buf1, fset, a)

	var buf2 bytes.Buffer
	*bp = &ast.ParenExpr{X: c}
	printer.Fprint(&buf2, fset, a)

	*bp = b

	// 3. Return whether they print the same (i.e., the parentheses are necessary).
	return buf1.String() == buf2.String()
}

func findExprField(a ast.Node, b ast.Expr) *ast.Expr {
	v := reflect.ValueOf(a).Elem()
	for i, n := 0, v.NumField(); i < n; i++ {
		// Interesting fields are either ast.Expr or []ast.Expr.
		switch f := v.Field(i).Interface().(type) {
		case ast.Expr:
			if f == b {
				return v.Field(i).Addr().Interface().(*ast.Expr)
			}
		case []ast.Expr:
			for i, e := range f {
				if e == b {
					return &f[i]
				}
			}
		}
	}
	log.Fatal("Failed to find b in a")
	return nil
}

func hasUntypedValue(n ast.Expr) bool {
	var v uvVisitor
	ast.Walk(&v, n)
	return v.found
}

type uvVisitor struct {
	found bool
}

func (v *uvVisitor) Visit(node ast.Node) ast.Visitor {
	// Short circuit.
	if v.found {
		return nil
	}

	switch node := node.(type) {
	case *ast.BinaryExpr:
		switch node.Op {
		case token.SHL, token.SHR, token.EQL, token.NEQ, token.LSS, token.GTR, token.LEQ, token.GEQ:
			// Shifts yield an untyped value if their LHS is untyped.
			// Comparisons yield an untyped boolean value.
			v.found = true
		}
	case *ast.Ident:
		if node.Name == "nil" {
			// Probably the universal untyped zero value.
			v.found = true
		}
	}

	if v.found {
		return nil
	}
	return v
}
<commit_msg>unconvert: tweak reflection usage<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/printer"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"sort"
	"sync"

	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
)

var fset = token.NewFileSet()

type Edit struct {
	Pos, End int
}

type editsByPos []Edit

func (e editsByPos) Len() int           { return len(e) }
func (e editsByPos) Less(i, j int) bool { return e[i].Pos < e[j].Pos }
func (e editsByPos) Swap(i, j int)      { e[i], e[j] = e[j], e[i] }

func apply(file string, edits []Edit) {
	if len(edits) == 0 {
		return
	}

	sort.Sort(editsByPos(edits))

	// Check for overlap.
	// TODO(mdempsky): Overlap can legally happen in bizarro expressions like
	// "(*[unsafe.Sizeof(int8(int8(0)))]byte)(*[1]byte)(nil)".
	for i := 1; i < len(edits); i++ {
		if edits[i-1].End > edits[i].Pos {
			log.Fatal("overlap")
		}
	}

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}

	n := edits[0].Pos
	for i := 1; i < len(edits); i++ {
		n += copy(buf[n:], buf[edits[i-1].End:edits[i].Pos])
	}
	n += copy(buf[n:], buf[edits[len(edits)-1].End:])
	buf = buf[:n]

	buf, err = format.Source(buf)
	if err != nil {
		log.Fatal(err)
	}

	err = ioutil.WriteFile(file, buf, 0)
	if err != nil {
		log.Fatal(err)
	}
}

var (
	flagAll   = flag.Bool("all", false, "type check all GOOS and GOARCH combinations")
	flagApply = flag.Bool("apply", false, "apply edits")
)

func main() {
	flag.Parse()

	var m map[string][]Edit
	if *flagAll {
		m = mergeEdits()
	} else {
		m = computeEdits(build.Default.GOOS, build.Default.GOARCH)
	}

	if *flagApply {
		for f, e := range m {
			apply(f, e)
		}
	} else {
		err := json.NewEncoder(os.Stdout).Encode(m)
		if err != nil {
			log.Fatal(err)
		}
		return
	}
}

var plats = [...]struct {
	goos, goarch string
}{
	{"linux", "386"},
	{"linux", "amd64"},
	{"linux", "arm"},
	{"linux", "arm64"},
	{"linux", "ppc64"},
	{"linux", "ppc64le"},
	{"nacl", "386"},
	{"nacl", "amd64p32"},
	{"nacl", "arm"},
	{"darwin", "386"},
	{"darwin", "amd64"},
	{"dragonfly", "amd64"},
	{"freebsd", "386"},
	{"freebsd", "amd64"},
	{"freebsd", "arm"},
	{"netbsd", "386"},
	{"netbsd", "amd64"},
	{"netbsd", "arm"},
	{"openbsd", "386"},
	{"openbsd", "amd64"},
	{"openbsd", "arm"},
	{"plan9", "386"},
	{"plan9", "amd64"},
	{"solaris", "amd64"},
	{"windows", "386"},
	{"windows", "amd64"},
}

func mergeEdits() map[string][]Edit {
	ch := make(chan map[string][]Edit)
	var wg sync.WaitGroup
	for _, plat := range plats {
		wg.Add(1)
		go func(goos, goarch string) {
			defer wg.Done()
			ch <- computeEdits(goos, goarch)
		}(plat.goos, plat.goarch)
	}
	go func() {
		wg.Wait()
		close(ch)
	}()

	m := make(map[string][]Edit)
	for m1 := range ch {
		for f, e := range m1 {
			if e0, ok := m[f]; ok {
				m[f] = intersect(e0, e)
			} else {
				m[f] = e
			}
		}
	}
	return m
}

func intersect(e1, e2 []Edit) []Edit {
	if len(e1) == 0 || len(e2) == 0 {
		return nil
	}

	set := make(map[Edit]bool, len(e1))
	for _, e := range e1 {
		set[e] = true
	}

	var res []Edit
	for _, e := range e2 {
		if set[e] {
			res = append(res, e)
		}
	}
	return res
}

func noImport(map[string]*types.Package, string) (*types.Package, error) {
	panic("go/loader said this wouldn't be called")
}

func computeEdits(os, arch string) map[string][]Edit {
	ctxt := build.Default
	ctxt.GOOS = os
	ctxt.GOARCH = arch
	ctxt.CgoEnabled = false

	var conf loader.Config
	conf.Fset = fset
	conf.Build = &ctxt
	conf.TypeChecker.Import = noImport
	for _, arg := range flag.Args() {
		conf.Import(arg)
	}
	prog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}

	type res struct {
		file  string
		edits []Edit
	}
	ch := make(chan res)
	var wg sync.WaitGroup
	for _, pkg := range prog.InitialPackages() {
		for _, file := range pkg.Files {
			wg.Add(1)
			go func(pkg *loader.PackageInfo, file *ast.File) {
				defer wg.Done()
				v := visitor{pkg: pkg, file: fset.File(file.Package)}
				ast.Walk(&v, file)
				sort.Sort(editsByPos(v.edits))
				ch <- res{v.file.Name(), v.edits}
			}(pkg, file)
		}
	}
	go func() {
		wg.Wait()
		close(ch)
	}()

	m := make(map[string][]Edit)
	for r := range ch {
		m[r.file] = r.edits
	}
	return m
}

type visitor struct {
	pkg   *loader.PackageInfo
	file  *token.File
	edits []Edit
	nodes []ast.Node
}

func (v *visitor) Visit(node ast.Node) ast.Visitor {
	if node != nil {
		v.nodes = append(v.nodes, node)
	} else {
		v.nodes = v.nodes[:len(v.nodes)-1]
	}
	if call, ok := node.(*ast.CallExpr); ok {
		v.unconvert(call)
	}
	return v
}

func (v *visitor) unconvert(call *ast.CallExpr) {
	// TODO(mdempsky): Handle useless multi-conversions.

	// Conversions have exactly one argument.
	if len(call.Args) != 1 || call.Ellipsis != token.NoPos {
		return
	}
	ft, ok := v.pkg.Types[call.Fun]
	if !ok {
		fmt.Println("Missing type for function")
		return
	}
	if !ft.IsType() {
		// Not a conversion.
		return
	}
	at, ok := v.pkg.Types[call.Args[0]]
	if !ok {
		fmt.Println("Missing type for argument")
	}
	if !types.Identical(ft.Type, at.Type) {
		// A real conversion.
		return
	}
	if at.Value != nil || hasUntypedValue(call.Args[0]) {
		// As a workaround for golang.org/issue/13061,
		// skip conversions that contain an untyped value.
		return
	}

	outer := v.nodes[len(v.nodes)-2]
	if keepParen(outer, call, call.Args[0]) {
		v.remove(call.Fun.Pos(), call.Lparen)
	} else {
		v.remove(call.Fun.Pos(), call.Lparen+1)
		v.remove(call.Rparen, call.Rparen+1)
	}
}

func (v *visitor) remove(pos, end token.Pos) {
	v.edits = append(v.edits, Edit{v.file.Offset(pos), v.file.Offset(end)})
}

func keepParen(a ast.Node, b, c ast.Expr) bool {
	// 1. Find the value in a that points to b.
	bp := findExprField(a, b)

	// 2. Try printing a with s/b/c/ and with s/b/(c)/.
	var buf1 bytes.Buffer
	*bp = c
	printer.Fprint(&buf1, fset, a)

	var buf2 bytes.Buffer
	*bp = &ast.ParenExpr{X: c}
	printer.Fprint(&buf2, fset, a)

	*bp = b

	// 3. Return whether they print the same (i.e., the parentheses are necessary).
	return buf1.String() == buf2.String()
}

func findExprField(a ast.Node, b ast.Expr) *ast.Expr {
	v := reflect.ValueOf(a).Elem()
	for i, n := 0, v.NumField(); i < n; i++ {
		// Interesting fields are either ast.Expr or []ast.Expr.
		switch f := v.Field(i).Addr().Interface().(type) {
		case *ast.Expr:
			if *f == b {
				return f
			}
		case *[]ast.Expr:
			for i, e := range *f {
				if e == b {
					return &(*f)[i]
				}
			}
		}
	}
	log.Fatal("Failed to find b in a")
	return nil
}

func hasUntypedValue(n ast.Expr) bool {
	var v uvVisitor
	ast.Walk(&v, n)
	return v.found
}

type uvVisitor struct {
	found bool
}

func (v *uvVisitor) Visit(node ast.Node) ast.Visitor {
	// Short circuit.
	if v.found {
		return nil
	}

	switch node := node.(type) {
	case *ast.BinaryExpr:
		switch node.Op {
		case token.SHL, token.SHR, token.EQL, token.NEQ, token.LSS, token.GTR, token.LEQ, token.GEQ:
			// Shifts yield an untyped value if their LHS is untyped.
			// Comparisons yield an untyped boolean value.
			v.found = true
		}
	case *ast.Ident:
		if node.Name == "nil" {
			// Probably the universal untyped zero value.
			v.found = true
		}
	}

	if v.found {
		return nil
	}
	return v
}
<|endoftext|>


<commit_before>package upnp

import (
	"sync"
	"testing"
)

// TestConcurrentUPNP tests that several threads calling Discover() concurrently
// succeed.
func TestConcurrentUPNP(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	// verify that a router exists
	_, err := Discover()
	if err != nil {
		t.Fatal(err)
	}

	// now try to concurrently Discover() using 20 threads
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		go func() {
			wg.Add(1)
			defer wg.Done()
			_, err := Discover()
			if err != nil {
				t.Fatal(err)
			}
		}()
	}
	wg.Wait()
}

func TestIGD(t *testing.T) {
	// connect to router
	d, err := Discover()
	if err != nil {
		t.Fatal(err)
	}

	// discover external IP
	ip, err := d.ExternalIP()
	if err != nil {
		t.Fatal(err)
	}
	t.Log("Your external IP is:", ip)

	// forward a port
	err = d.Forward(9001, "upnp test")
	if err != nil {
		t.Fatal(err)
	}

	// un-forward a port
	err = d.Clear(9001)
	if err != nil {
		t.Fatal(err)
	}

	// record router's location
	loc := d.Location()
	if err != nil {
		t.Fatal(err)
	}

	// connect to router directly
	d, err = Load(loc)
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>skip tests which require a router if not discovered<commit_after>package upnp

import (
	"sync"
	"testing"
)

// TestConcurrentUPNP tests that several threads calling Discover() concurrently
// succeed.
func TestConcurrentUPNP(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	// verify that a router exists
	_, err := Discover()
	if err != nil {
		t.Skip(err)
	}

	// now try to concurrently Discover() using 20 threads
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		go func() {
			wg.Add(1)
			defer wg.Done()
			_, err := Discover()
			if err != nil {
				t.Fatal(err)
			}
		}()
	}
	wg.Wait()
}

func TestIGD(t *testing.T) {
	// connect to router
	d, err := Discover()
	if err != nil {
		t.Skip(err)
	}

	// discover external IP
	ip, err := d.ExternalIP()
	if err != nil {
		t.Fatal(err)
	}
	t.Log("Your external IP is:", ip)

	// forward a port
	err = d.Forward(9001, "upnp test")
	if err != nil {
		t.Fatal(err)
	}

	// un-forward a port
	err = d.Clear(9001)
	if err != nil {
		t.Fatal(err)
	}

	// record router's location
	loc := d.Location()
	if err != nil {
		t.Fatal(err)
	}

	// connect to router directly
	d, err = Load(loc)
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package util

import (
	"container/list"
	"reflect"
)

// Channel that does not block when items are sent. To use the struct, simply
// send on the Send channel and receive on the Recv channel. Items will be
// stored internally until they are received. Closing the Send channel will
// cause the Recv channel to be also be closed once all items are received.
type NonBlockingChan struct {
	Send chan<- interface{}
	Recv <-chan interface{}
}

// Create a new non-blocking channel.
func NewNonBlockingChan() *NonBlockingChan {

	// Create the two channels that will be used for sending and receiving
	var (
		send = make(chan interface{})
		recv = make(chan interface{})
	)

	// Assign the channels to the public members of the struct (which limits
	// each of their direction)
	n := &NonBlockingChan{
		Send: send,
		Recv: recv,
	}

	// Start a goroutine to perform the sending and receiving
	go func() {

		// Create the list that will temporarily hold items for receiving and
		// create a boolean to track whether the Send channel has been closed
		var (
			items      = list.New()
			sendClosed = false
		)

		// Create constants for the select case array
		const (
			incomingCase = iota
			outgoingCase
			numCases
		)

		for {

			// Close the Recv channel and quit if the Send channel was closed
			// and there are no more items left in the list to send
			if sendClosed && items.Len() == 0 {
				close(recv)
				break
			}

			// Because the number of cases changes depending on the current
			// state, a dynamic select must be used
			cases := make([]reflect.SelectCase, numCases)
			cases[incomingCase].Dir = reflect.SelectRecv
			cases[outgoingCase].Dir = reflect.SelectSend

			// If the Send channel is still open, then add a select case to
			// receive from it - otherwise, leave it empty
			if !sendClosed {
				cases[incomingCase].Chan = reflect.ValueOf(send)
			}

			// If the list contains at least one item, add a select case to
			// send the first item from the list on the Recv channel
			if items.Len() > 0 {
				cases[outgoingCase].Chan = reflect.ValueOf(recv)
				cases[outgoingCase].Send = reflect.ValueOf(items.Front().Value)
			}

			// Three items will be returned - the index of the select case that
			// was chosen, the item received (if applicable), and the state of
			// the channel (open or closed)
			i, item, ok := reflect.Select(cases)

			// Switch on the chosen channel
			switch i {
			case incomingCase:

				// If an item was received, add it to the list - otherwise,
				// make a note that the Send channel was closed
				if ok {
					items.PushBack(item.Interface())
				} else {
					sendClosed = true
				}
			case outgoingCase:

				// The first item was sent, so it can be removed from the list
				items.Remove(items.Front())
			}
		}
	}()

	return n
}
<commit_msg>Refactored NonBlockingChan - reflect is no longer needed.<commit_after>package util

import (
	"container/list"
)

// Channel that does not block when items are sent. To use the type, simply
// send on the Send channel and receive on the Recv channel. Items will be
// stored internally until they are received. Closing the Send channel will
// cause the Recv channel to be also be closed after all items are received.
type NonBlockingChan struct {
	Send chan<- interface{}
	Recv <-chan interface{}
}

// Create a new non-blocking channel.
func NewNonBlockingChan() *NonBlockingChan {
	var (
		in  = make(chan interface{})
		out = make(chan interface{})
		n   = &NonBlockingChan{
			Send: in,
			Recv: out,
		}
	)
	go func() {
		var (
			items    = list.New()
			inClosed = false
		)
		for {
			if inClosed && items.Len() == 0 {
				close(out)
				break
			}
			var (
				inChan  chan interface{}
				outChan chan interface{}
				outVal  interface{}
			)
			if !inClosed {
				inChan = in
			}
			if items.Len() > 0 {
				outChan, outVal = out, items.Front().Value
			}
			select {
			case i, ok := <-inChan:
				if ok {
					items.PushBack(i)
				} else {
					inClosed = true
				}
			case outChan <- outVal:
				items.Remove(items.Front())
			}
		}
	}()
	return n
}
<|endoftext|>


<commit_before>package kcp

import (
	"encoding/binary"
	"log"
	"sync"

	"github.com/klauspost/reedsolomon"
)

const (
	fecHeaderSize      = 6
	fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size
	typeData           = 0xf1
	typeFEC            = 0xf2
	fecExpire          = 30000 // 30s
)

type (
	// FEC defines forward error correction for packets
	FEC struct {
		rx           []fecPacket // ordered receive queue
		rxlimit      int         // queue size limit
		dataShards   int
		parityShards int
		shardSize    int
		next         uint32 // next seqid
		enc          reedsolomon.Encoder
		shards       [][]byte
		shardsflag   []bool
		paws         uint32 // Protect Against Wrapped Sequence numbers
		lastCheck    uint32
		xmitBuf      sync.Pool
	}

	fecPacket struct {
		seqid uint32
		flag  uint16
		data  []byte
		ts    uint32
	}
)

func newFEC(rxlimit, dataShards, parityShards int) *FEC {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	if rxlimit < dataShards+parityShards {
		return nil
	}

	fec := new(FEC)
	fec.rxlimit = rxlimit
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	fec.paws = (0xffffffff/uint32(fec.shardSize) - 1) * uint32(fec.shardSize)
	enc, err := reedsolomon.New(dataShards, parityShards)
	if err != nil {
		log.Println(err)
		return nil
	}
	fec.enc = enc
	fec.shards = make([][]byte, fec.shardSize)
	fec.shardsflag = make([]bool, fec.shardSize)
	fec.xmitBuf.New = func() interface{} {
		return make([]byte, mtuLimit)
	}

	return fec
}

// decode a fec packet
func (fec *FEC) decode(data []byte) fecPacket {
	var pkt fecPacket
	pkt.seqid = binary.LittleEndian.Uint32(data)
	pkt.flag = binary.LittleEndian.Uint16(data[4:])
	pkt.ts = currentMs()
	// allocate memory & copy
	buf := fec.xmitBuf.Get().([]byte)
	copy(buf, data[6:])
	n := len(data[6:])
	xorBytes(buf[n:], buf[n:], buf[n:])
	pkt.data = buf
	return pkt
}

func (fec *FEC) markData(data []byte) {
	binary.LittleEndian.PutUint32(data, fec.next)
	binary.LittleEndian.PutUint16(data[4:], typeData)
	fec.next++
}

func (fec *FEC) markFEC(data []byte) {
	binary.LittleEndian.PutUint32(data, fec.next)
	binary.LittleEndian.PutUint16(data[4:], typeFEC)
	fec.next++
	if fec.next >= fec.paws { // paws would only occurs in markFEC
		fec.next = 0
	}
}

// input a fec packet
func (fec *FEC) input(pkt fecPacket) (recovered [][]byte) {
	// expiration
	now := currentMs()
	if now-fec.lastCheck >= fecExpire {
		var rx []fecPacket
		for k := range fec.rx {
			if now-fec.rx[k].ts < fecExpire {
				rx = append(rx, fec.rx[k])
			} else {
				fec.xmitBuf.Put(fec.rx[k].data)
			}
		}
		fec.rx = rx
		fec.lastCheck = now
	}

	// insertion
	n := len(fec.rx) - 1
	insertIdx := 0
	for i := n; i >= 0; i-- {
		if pkt.seqid == fec.rx[i].seqid { // de-duplicate
			fec.xmitBuf.Put(pkt.data)
			return nil
		} else if pkt.seqid > fec.rx[i].seqid { // insertion
			insertIdx = i + 1
			break
		}
	}

	// insert into ordered rx queue
	if insertIdx == n+1 {
		fec.rx = append(fec.rx, pkt)
	} else {
		fec.rx = append(fec.rx, fecPacket{})
		copy(fec.rx[insertIdx+1:], fec.rx[insertIdx:])
		fec.rx[insertIdx] = pkt
	}

	// shard range for current packet
	shardBegin := pkt.seqid - pkt.seqid%uint32(fec.shardSize)
	shardEnd := shardBegin + uint32(fec.shardSize) - 1

	// max search range in ordered queue for current shard
	searchBegin := insertIdx - int(pkt.seqid%uint32(fec.shardSize))
	if searchBegin < 0 {
		searchBegin = 0
	}
	searchEnd := searchBegin + fec.shardSize - 1
	if searchEnd >= len(fec.rx) {
		searchEnd = len(fec.rx) - 1
	}

	if searchEnd > searchBegin && searchEnd-searchBegin+1 >= fec.dataShards {
		numshard := 0
		numDataShard := 0
		first := -1
		maxlen := 0
		shards := fec.shards
		shardsflag := fec.shardsflag
		for k := range fec.shards {
			shards[k] = nil
			shardsflag[k] = false
		}

		for i := searchBegin; i <= searchEnd; i++ {
			seqid := fec.rx[i].seqid
			if seqid > shardEnd {
				break
			} else if seqid >= shardBegin {
				shards[seqid%uint32(fec.shardSize)] = fec.rx[i].data
				shardsflag[seqid%uint32(fec.shardSize)] = true
				numshard++
				if fec.rx[i].flag == typeData {
					numDataShard++
				}
				if numshard == 1 {
					first = i
				}
				if len(fec.rx[i].data) > maxlen {
					maxlen = len(fec.rx[i].data)
				}
			}
		}

		if numDataShard == fec.dataShards { // no lost
			for i := first; i < first+numshard; i++ { // free
				fec.xmitBuf.Put(fec.rx[i].data)
			}
			copy(fec.rx[first:], fec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				fec.rx[len(fec.rx)-1-i] = fecPacket{}
			}
			fec.rx = fec.rx[:len(fec.rx)-numshard]
		} else if numshard >= fec.dataShards { // recoverable
			for k := range shards {
				if shards[k] != nil {
					shards[k] = shards[k][:maxlen]
				}
			}
			if err := fec.enc.Reconstruct(shards); err == nil {
				for k := range shards[:fec.dataShards] {
					if !shardsflag[k] {
						recovered = append(recovered, shards[k])
					}
				}
			} else {
				log.Println(err)
			}

			for i := first; i < first+numshard; i++ { // free
				fec.xmitBuf.Put(fec.rx[i].data)
			}
			copy(fec.rx[first:], fec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				fec.rx[len(fec.rx)-1-i] = fecPacket{}
			}
			fec.rx = fec.rx[:len(fec.rx)-numshard]
		}
	}

	// keep rxlimit
	if len(fec.rx) > fec.rxlimit {
		fec.xmitBuf.Put(fec.rx[0].data) // free
		fec.rx[0].data = nil
		fec.rx = fec.rx[1:]
	}
	return
}

func (fec *FEC) calcECC(data [][]byte, offset, maxlen int) (ecc [][]byte) {
	if len(data) != fec.shardSize {
		log.Println("mismatch", len(data), fec.shardSize)
		return nil
	}
	shards := make([][]byte, fec.shardSize)
	for k := range shards {
		shards[k] = data[k][offset:maxlen]
	}

	if err := fec.enc.Encode(shards); err != nil {
		log.Println(err)
		return nil
	}
	return data[fec.dataShards:]
}
<commit_msg>fix a stupid code<commit_after>package kcp

import (
	"encoding/binary"
	"log"
	"sync"

	"github.com/klauspost/reedsolomon"
)

const (
	fecHeaderSize      = 6
	fecHeaderSizePlus2 = fecHeaderSize + 2 // plus 2B data size
	typeData           = 0xf1
	typeFEC            = 0xf2
	fecExpire          = 30000 // 30s
)

type (
	// FEC defines forward error correction for packets
	FEC struct {
		rx           []fecPacket // ordered receive queue
		rxlimit      int         // queue size limit
		dataShards   int
		parityShards int
		shardSize    int
		next         uint32 // next seqid
		enc          reedsolomon.Encoder
		shards       [][]byte
		shardsflag   []bool
		paws         uint32 // Protect Against Wrapped Sequence numbers
		lastCheck    uint32
		xmitBuf      sync.Pool
	}

	fecPacket struct {
		seqid uint32
		flag  uint16
		data  []byte
		ts    uint32
	}
)

func newFEC(rxlimit, dataShards, parityShards int) *FEC {
	if dataShards <= 0 || parityShards <= 0 {
		return nil
	}
	if rxlimit < dataShards+parityShards {
		return nil
	}

	fec := new(FEC)
	fec.rxlimit = rxlimit
	fec.dataShards = dataShards
	fec.parityShards = parityShards
	fec.shardSize = dataShards + parityShards
	fec.paws = (0xffffffff/uint32(fec.shardSize) - 1) * uint32(fec.shardSize)
	enc, err := reedsolomon.New(dataShards, parityShards)
	if err != nil {
		log.Println(err)
		return nil
	}
	fec.enc = enc
	fec.shards = make([][]byte, fec.shardSize)
	fec.shardsflag = make([]bool, fec.shardSize)
	fec.xmitBuf.New = func() interface{} {
		return make([]byte, mtuLimit)
	}

	return fec
}

// decode a fec packet
func (fec *FEC) decode(data []byte) fecPacket {
	var pkt fecPacket
	pkt.seqid = binary.LittleEndian.Uint32(data)
	pkt.flag = binary.LittleEndian.Uint16(data[4:])
	pkt.ts = currentMs()
	// allocate memory & copy
	buf := fec.xmitBuf.Get().([]byte)
	n := copy(buf, data[6:])
	xorBytes(buf[n:], buf[n:], buf[n:])
	pkt.data = buf
	return pkt
}

func (fec *FEC) markData(data []byte) {
	binary.LittleEndian.PutUint32(data, fec.next)
	binary.LittleEndian.PutUint16(data[4:], typeData)
	fec.next++
}

func (fec *FEC) markFEC(data []byte) {
	binary.LittleEndian.PutUint32(data, fec.next)
	binary.LittleEndian.PutUint16(data[4:], typeFEC)
	fec.next++
	if fec.next >= fec.paws { // paws would only occurs in markFEC
		fec.next = 0
	}
}

// input a fec packet
func (fec *FEC) input(pkt fecPacket) (recovered [][]byte) {
	// expiration
	now := currentMs()
	if now-fec.lastCheck >= fecExpire {
		var rx []fecPacket
		for k := range fec.rx {
			if now-fec.rx[k].ts < fecExpire {
				rx = append(rx, fec.rx[k])
			} else {
				fec.xmitBuf.Put(fec.rx[k].data)
			}
		}
		fec.rx = rx
		fec.lastCheck = now
	}

	// insertion
	n := len(fec.rx) - 1
	insertIdx := 0
	for i := n; i >= 0; i-- {
		if pkt.seqid == fec.rx[i].seqid { // de-duplicate
			fec.xmitBuf.Put(pkt.data)
			return nil
		} else if pkt.seqid > fec.rx[i].seqid { // insertion
			insertIdx = i + 1
			break
		}
	}

	// insert into ordered rx queue
	if insertIdx == n+1 {
		fec.rx = append(fec.rx, pkt)
	} else {
		fec.rx = append(fec.rx, fecPacket{})
		copy(fec.rx[insertIdx+1:], fec.rx[insertIdx:])
		fec.rx[insertIdx] = pkt
	}

	// shard range for current packet
	shardBegin := pkt.seqid - pkt.seqid%uint32(fec.shardSize)
	shardEnd := shardBegin + uint32(fec.shardSize) - 1

	// max search range in ordered queue for current shard
	searchBegin := insertIdx - int(pkt.seqid%uint32(fec.shardSize))
	if searchBegin < 0 {
		searchBegin = 0
	}
	searchEnd := searchBegin + fec.shardSize - 1
	if searchEnd >= len(fec.rx) {
		searchEnd = len(fec.rx) - 1
	}

	if searchEnd > searchBegin && searchEnd-searchBegin+1 >= fec.dataShards {
		numshard := 0
		numDataShard := 0
		first := -1
		maxlen := 0
		shards := fec.shards
		shardsflag := fec.shardsflag
		for k := range fec.shards {
			shards[k] = nil
			shardsflag[k] = false
		}

		for i := searchBegin; i <= searchEnd; i++ {
			seqid := fec.rx[i].seqid
			if seqid > shardEnd {
				break
			} else if seqid >= shardBegin {
				shards[seqid%uint32(fec.shardSize)] = fec.rx[i].data
				shardsflag[seqid%uint32(fec.shardSize)] = true
				numshard++
				if fec.rx[i].flag == typeData {
					numDataShard++
				}
				if numshard == 1 {
					first = i
				}
				if len(fec.rx[i].data) > maxlen {
					maxlen = len(fec.rx[i].data)
				}
			}
		}

		if numDataShard == fec.dataShards { // no lost
			for i := first; i < first+numshard; i++ { // free
				fec.xmitBuf.Put(fec.rx[i].data)
			}
			copy(fec.rx[first:], fec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				fec.rx[len(fec.rx)-1-i] = fecPacket{}
			}
			fec.rx = fec.rx[:len(fec.rx)-numshard]
		} else if numshard >= fec.dataShards { // recoverable
			for k := range shards {
				if shards[k] != nil {
					shards[k] = shards[k][:maxlen]
				}
			}
			if err := fec.enc.Reconstruct(shards); err == nil {
				for k := range shards[:fec.dataShards] {
					if !shardsflag[k] {
						recovered = append(recovered, shards[k])
					}
				}
			} else {
				log.Println(err)
			}

			for i := first; i < first+numshard; i++ { // free
				fec.xmitBuf.Put(fec.rx[i].data)
			}
			copy(fec.rx[first:], fec.rx[first+numshard:])
			for i := 0; i < numshard; i++ { // dereference
				fec.rx[len(fec.rx)-1-i] = fecPacket{}
			}
			fec.rx = fec.rx[:len(fec.rx)-numshard]
		}
	}

	// keep rxlimit
	if len(fec.rx) > fec.rxlimit {
		fec.xmitBuf.Put(fec.rx[0].data) // free
		fec.rx[0].data = nil
		fec.rx = fec.rx[1:]
	}
	return
}

func (fec *FEC) calcECC(data [][]byte, offset, maxlen int) (ecc [][]byte) {
	if len(data) != fec.shardSize {
		log.Println("mismatch", len(data), fec.shardSize)
		return nil
	}
	shards := make([][]byte, fec.shardSize)
	for k := range shards {
		shards[k] = data[k][offset:maxlen]
	}

	if err := fec.enc.Encode(shards); err != nil {
		log.Println(err)
		return nil
	}
	return data[fec.dataShards:]
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package upgrade

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
)

const (
	validConfig = `apiVersion: kubeadm.k8s.io/v1beta1
kind: ClusterConfiguration
kubernetesVersion: 1.13.0
`
)

func TestGetK8sVersionFromUserInput(t *testing.T) {
	var tcases = []struct {
		name               string
		isVersionMandatory bool
		clusterConfig      string
		args               []string
		expectedErr        bool
		expectedVersion    string
	}{
		{
			name:               "No config and version as an argument",
			isVersionMandatory: true,
			args:               []string{"v1.13.1"},
			expectedVersion:    "v1.13.1",
		},
		{
			name:               "Neither config nor version specified",
			isVersionMandatory: true,
			expectedErr:        true,
		},
		{
			name:               "No config and empty version as an argument",
			isVersionMandatory: true,
			args:               []string{""},
			expectedErr:        true,
		},
		{
			name:               "Valid config, but no version specified",
			isVersionMandatory: true,
			clusterConfig:      validConfig,
			expectedVersion:    "v1.13.0",
		},
		{
			name:               "Valid config and different version specified",
			isVersionMandatory: true,
			clusterConfig:      validConfig,
			args:               []string{"v1.13.1"},
			expectedVersion:    "v1.13.1",
		},
		{
			name: "Version is optional",
		},
	}
	for tnum, tt := range tcases {
		t.Run(tt.name, func(t *testing.T) {
			flags := &applyPlanFlags{}
			if len(tt.clusterConfig) > 0 {
				tmpfile := fmt.Sprintf("/tmp/kubeadm-upgrade-common-test-%d-%d.yaml", tnum, time.Now().Unix())
				if err := ioutil.WriteFile(tmpfile, []byte(tt.clusterConfig), 0666); err != nil {
					t.Fatalf("Failed to create test config file: %+v", err)
				}
				defer os.Remove(tmpfile)

				flags.cfgPath = tmpfile
			}

			userVersion, err := getK8sVersionFromUserInput(flags, tt.args, tt.isVersionMandatory)

			if err == nil && tt.expectedErr {
				t.Error("Expected error, but got success")
			}
			if err != nil && !tt.expectedErr {
				t.Errorf("Unexpected error: %+v", err)
			}
			if userVersion != tt.expectedVersion {
				t.Errorf("Expected '%s', but got '%s'", tt.expectedVersion, userVersion)
			}
		})
	}
}

func TestEnforceRequirements(t *testing.T) {
	tcases := []struct {
		name          string
		newK8sVersion string
		dryRun        bool
		flags         applyPlanFlags
		expectedErr   bool
	}{
		{
			name:        "Fail pre-flight check",
			expectedErr: true,
		},
		{
			name: "Bogus preflight check disabled when also 'all' is specified",
			flags: applyPlanFlags{
				ignorePreflightErrors: []string{"bogusvalue", "all"},
			},
			expectedErr: true,
		},
		{
			name: "Fail to create client",
			flags: applyPlanFlags{
				ignorePreflightErrors: []string{"all"},
			},
			expectedErr: true,
		},
	}
	for _, tt := range tcases {
		t.Run(tt.name, func(t *testing.T) {
			_, _, _, err := enforceRequirements(&tt.flags, tt.dryRun, tt.newK8sVersion)

			if err == nil && tt.expectedErr {
				t.Error("Expected error, but got success")
			}
			if err != nil && !tt.expectedErr {
				t.Errorf("Unexpected error: %+v", err)
			}
		})
	}
}

func TestPrintConfiguration(t *testing.T) {
	var tests = []struct {
		name          string
		cfg           *kubeadmapi.ClusterConfiguration
		buf           *bytes.Buffer
		expectedBytes []byte
	}{
		{
			name:          "config is nil",
			cfg:           nil,
			expectedBytes: []byte(""),
		},
		{
			name: "cluster config with local Etcd",
			cfg: &kubeadmapi.ClusterConfiguration{
				KubernetesVersion: "v1.7.1",
				Etcd: kubeadmapi.Etcd{
					Local: &kubeadmapi.LocalEtcd{
						DataDir: "/some/path",
					},
				},
				DNS: kubeadmapi.DNS{
					Type: kubeadmapi.CoreDNS,
				},
			},
			expectedBytes: []byte(`[upgrade/config] Configuration used:
	apiServer: {}
	apiVersion: kubeadm.k8s.io/v1beta1
	certificatesDir: ""
	controlPlaneEndpoint: ""
	controllerManager: {}
	dns:
	  type: CoreDNS
	etcd:
	  local:
	    dataDir: /some/path
	imageRepository: ""
	kind: ClusterConfiguration
	kubernetesVersion: v1.7.1
	networking:
	  dnsDomain: ""
	  podSubnet: ""
	  serviceSubnet: ""
	scheduler: {}
`),
		},
		{
			name: "cluster config with ServiceSubnet and external Etcd",
			cfg: &kubeadmapi.ClusterConfiguration{
				KubernetesVersion: "v1.7.1",
				Networking: kubeadmapi.Networking{
					ServiceSubnet: "10.96.0.1/12",
				},
				Etcd: kubeadmapi.Etcd{
					External: &kubeadmapi.ExternalEtcd{
						Endpoints: []string{"https://one-etcd-instance:2379"},
					},
				},
				DNS: kubeadmapi.DNS{
					Type: kubeadmapi.CoreDNS,
				},
			},
			expectedBytes: []byte(`[upgrade/config] Configuration used:
	apiServer: {}
	apiVersion: kubeadm.k8s.io/v1beta1
	certificatesDir: ""
	controlPlaneEndpoint: ""
	controllerManager: {}
	dns:
	  type: CoreDNS
	etcd:
	  external:
	    caFile: ""
	    certFile: ""
	    endpoints:
	    - https://one-etcd-instance:2379
	    keyFile: ""
	imageRepository: ""
	kind: ClusterConfiguration
	kubernetesVersion: v1.7.1
	networking:
	  dnsDomain: ""
	  podSubnet: ""
	  serviceSubnet: 10.96.0.1/12
	scheduler: {}
`),
		},
	}
	for _, rt := range tests {
		t.Run(rt.name, func(t *testing.T) {
			rt.buf = bytes.NewBufferString("")
			printConfiguration(rt.cfg, rt.buf)
			actualBytes := rt.buf.Bytes()
			if !bytes.Equal(actualBytes, rt.expectedBytes) {
				t.Errorf(
					"failed PrintConfiguration:\n\texpected: %q\n\t  actual: %q",
					string(rt.expectedBytes),
					string(actualBytes),
				)
			}
		})
	}
}
<commit_msg>kubeadm: Don't hardcode temp path in a test<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package upgrade

import (
	"bytes"
	"io/ioutil"
	"os"
	"testing"

	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
)

const (
	validConfig = `apiVersion: kubeadm.k8s.io/v1beta1
kind: ClusterConfiguration
kubernetesVersion: 1.13.0
`
)

func TestGetK8sVersionFromUserInput(t *testing.T) {
	var tcases = []struct {
		name               string
		isVersionMandatory bool
		clusterConfig      string
		args               []string
		expectedErr        bool
		expectedVersion    string
	}{
		{
			name:               "No config and version as an argument",
			isVersionMandatory: true,
			args:               []string{"v1.13.1"},
			expectedVersion:    "v1.13.1",
		},
		{
			name:               "Neither config nor version specified",
			isVersionMandatory: true,
			expectedErr:        true,
		},
		{
			name:               "No config and empty version as an argument",
			isVersionMandatory: true,
			args:               []string{""},
			expectedErr:        true,
		},
		{
			name:               "Valid config, but no version specified",
			isVersionMandatory: true,
			clusterConfig:      validConfig,
			expectedVersion:    "v1.13.0",
		},
		{
			name:               "Valid config and different version specified",
			isVersionMandatory: true,
			clusterConfig:      validConfig,
			args:               []string{"v1.13.1"},
			expectedVersion:    "v1.13.1",
		},
		{
			name: "Version is optional",
		},
	}
	for _, tt := range tcases {
		t.Run(tt.name, func(t *testing.T) {
			flags := &applyPlanFlags{}
			if len(tt.clusterConfig) > 0 {
				file, err := ioutil.TempFile("", "kubeadm-upgrade-common-test-*.yaml")
				if err != nil {
					t.Fatalf("Failed to create test config file: %+v", err)
				}

				tmpFileName := file.Name()
				defer os.Remove(tmpFileName)

				_, err = file.WriteString(tt.clusterConfig)
				file.Close()
				if err != nil {
					t.Fatalf("Failed to write test config file contents: %+v", err)
				}

				flags.cfgPath = tmpFileName
			}

			userVersion, err := getK8sVersionFromUserInput(flags, tt.args, tt.isVersionMandatory)

			if err == nil && tt.expectedErr {
				t.Error("Expected error, but got success")
			}
			if err != nil && !tt.expectedErr {
				t.Errorf("Unexpected error: %+v", err)
			}
			if userVersion != tt.expectedVersion {
				t.Errorf("Expected '%s', but got '%s'", tt.expectedVersion, userVersion)
			}
		})
	}
}

func TestEnforceRequirements(t *testing.T) {
	tcases := []struct {
		name          string
		newK8sVersion string
		dryRun        bool
		flags         applyPlanFlags
		expectedErr   bool
	}{
		{
			name:        "Fail pre-flight check",
			expectedErr: true,
		},
		{
			name: "Bogus preflight check disabled when also 'all' is specified",
			flags: applyPlanFlags{
				ignorePreflightErrors: []string{"bogusvalue", "all"},
			},
			expectedErr: true,
		},
		{
			name: "Fail to create client",
			flags: applyPlanFlags{
				ignorePreflightErrors: []string{"all"},
			},
			expectedErr: true,
		},
	}
	for _, tt := range tcases {
		t.Run(tt.name, func(t *testing.T) {
			_, _, _, err := enforceRequirements(&tt.flags, tt.dryRun, tt.newK8sVersion)

			if err == nil && tt.expectedErr {
				t.Error("Expected error, but got success")
			}
			if err != nil && !tt.expectedErr {
				t.Errorf("Unexpected error: %+v", err)
			}
		})
	}
}

func TestPrintConfiguration(t *testing.T) {
	var tests = []struct {
		name          string
		cfg           *kubeadmapi.ClusterConfiguration
		buf           *bytes.Buffer
		expectedBytes []byte
	}{
		{
			name:          "config is nil",
			cfg:           nil,
			expectedBytes: []byte(""),
		},
		{
			name: "cluster config with local Etcd",
			cfg: &kubeadmapi.ClusterConfiguration{
				KubernetesVersion: "v1.7.1",
				Etcd: kubeadmapi.Etcd{
					Local: &kubeadmapi.LocalEtcd{
						DataDir: "/some/path",
					},
				},
				DNS: kubeadmapi.DNS{
					Type: kubeadmapi.CoreDNS,
				},
			},
			expectedBytes: []byte(`[upgrade/config] Configuration used:
	apiServer: {}
	apiVersion: kubeadm.k8s.io/v1beta1
	certificatesDir: ""
	controlPlaneEndpoint: ""
	controllerManager: {}
	dns:
	  type: CoreDNS
	etcd:
	  local:
	    dataDir: /some/path
	imageRepository: ""
	kind: ClusterConfiguration
	kubernetesVersion: v1.7.1
	networking:
	  dnsDomain: ""
	  podSubnet: ""
	  serviceSubnet: ""
	scheduler: {}
`),
		},
		{
			name: "cluster config with ServiceSubnet and external Etcd",
			cfg: &kubeadmapi.ClusterConfiguration{
				KubernetesVersion: "v1.7.1",
				Networking: kubeadmapi.Networking{
					ServiceSubnet: "10.96.0.1/12",
				},
				Etcd: kubeadmapi.Etcd{
					External: &kubeadmapi.ExternalEtcd{
						Endpoints: []string{"https://one-etcd-instance:2379"},
					},
				},
				DNS: kubeadmapi.DNS{
					Type: kubeadmapi.CoreDNS,
				},
			},
			expectedBytes: []byte(`[upgrade/config] Configuration used:
	apiServer: {}
	apiVersion: kubeadm.k8s.io/v1beta1
	certificatesDir: ""
	controlPlaneEndpoint: ""
	controllerManager: {}
	dns:
	  type: CoreDNS
	etcd:
	  external:
	    caFile: ""
	    certFile: ""
	    endpoints:
	    - https://one-etcd-instance:2379
	    keyFile: ""
	imageRepository: ""
	kind: ClusterConfiguration
	kubernetesVersion: v1.7.1
	networking:
	  dnsDomain: ""
	  podSubnet: ""
	  serviceSubnet: 10.96.0.1/12
	scheduler: {}
`),
		},
	}
	for _, rt := range tests {
		t.Run(rt.name, func(t *testing.T) {
			rt.buf = bytes.NewBufferString("")
			printConfiguration(rt.cfg, rt.buf)
			actualBytes := rt.buf.Bytes()
			if !bytes.Equal(actualBytes, rt.expectedBytes) {
				t.Errorf(
					"failed PrintConfiguration:\n\texpected: %q\n\t  actual: %q",
					string(rt.expectedBytes),
					string(actualBytes),
				)
			}
		})
	}
}
<|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 certs

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"testing"

	certutil "k8s.io/client-go/util/cert"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
)

func TestCreatePKIAssets(t *testing.T) {
	tmpdir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("Couldn't create tmpdir")
	}
	defer os.Remove(tmpdir)

	var tests = []struct {
		cfg      *kubeadmapi.MasterConfiguration
		expected bool
	}{
		{
			cfg:      &kubeadmapi.MasterConfiguration{},
			expected: false,
		},
		{
			// CIDR too small
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "10.0.0.1/1"},
			},
			expected: false,
		},
		{
			// CIDR invalid
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "invalid"},
			},
			expected: false,
		},
		{
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "10.0.0.1/24"},
			},
			expected: true,
		},
	}
	for _, rt := range tests {
		actual := CreatePKIAssets(rt.cfg, fmt.Sprintf("%s/etc/kubernetes/pki", tmpdir))
		if (actual == nil) != rt.expected {
			t.Errorf(
				"failed CreatePKIAssets with an error:\n\texpected: %t\n\t  actual: %t",
				rt.expected,
				(actual == nil),
			)
		}
	}
}

func TestCheckAltNamesExist(t *testing.T) {
	var tests = []struct {
		IPs              []net.IP
		DNSNames         []string
		requiredAltNames certutil.AltNames
		succeed          bool
	}{
		{
			// equal
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          true,
		},
		{
			// the loaded cert has more ips than required, ok
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("192.168.2.5"), net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"a", "foo", "b", "bar", "baz"},
			succeed:          true,
		},
		{
			// the loaded cert doesn't have all ips
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.2.5"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          false,
		},
		{
			// the loaded cert doesn't have all ips
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "b", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          false,
		},
	}

	for _, rt := range tests {
		succeeded := checkAltNamesExist(rt.IPs, rt.DNSNames, rt.requiredAltNames)
		if succeeded != rt.succeed {
			t.Errorf(
				"failed checkAltNamesExist:\n\texpected: %t\n\t  actual: %t",
				rt.succeed,
				succeeded,
			)
		}
	}
}
<commit_msg>kubeadm: change os.Remove to os.RemoveAll<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 certs

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"testing"

	certutil "k8s.io/client-go/util/cert"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm"
)

func TestCreatePKIAssets(t *testing.T) {
	tmpdir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("Couldn't create tmpdir")
	}
	defer os.RemoveAll(tmpdir)

	var tests = []struct {
		cfg      *kubeadmapi.MasterConfiguration
		expected bool
	}{
		{
			cfg:      &kubeadmapi.MasterConfiguration{},
			expected: false,
		},
		{
			// CIDR too small
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "10.0.0.1/1"},
			},
			expected: false,
		},
		{
			// CIDR invalid
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "invalid"},
			},
			expected: false,
		},
		{
			cfg: &kubeadmapi.MasterConfiguration{
				API:        kubeadmapi.API{AdvertiseAddresses: []string{"10.0.0.1"}},
				Networking: kubeadmapi.Networking{ServiceSubnet: "10.0.0.1/24"},
			},
			expected: true,
		},
	}
	for _, rt := range tests {
		actual := CreatePKIAssets(rt.cfg, fmt.Sprintf("%s/etc/kubernetes/pki", tmpdir))
		if (actual == nil) != rt.expected {
			t.Errorf(
				"failed CreatePKIAssets with an error:\n\texpected: %t\n\t  actual: %t",
				rt.expected,
				(actual == nil),
			)
		}
	}
}

func TestCheckAltNamesExist(t *testing.T) {
	var tests = []struct {
		IPs              []net.IP
		DNSNames         []string
		requiredAltNames certutil.AltNames
		succeed          bool
	}{
		{
			// equal
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          true,
		},
		{
			// the loaded cert has more ips than required, ok
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("192.168.2.5"), net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"a", "foo", "b", "bar", "baz"},
			succeed:          true,
		},
		{
			// the loaded cert doesn't have all ips
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.2.5"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          false,
		},
		{
			// the loaded cert doesn't have all ips
			requiredAltNames: certutil.AltNames{IPs: []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")}, DNSNames: []string{"foo", "bar", "b", "baz"}},
			IPs:              []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("192.168.1.2")},
			DNSNames:         []string{"foo", "bar", "baz"},
			succeed:          false,
		},
	}

	for _, rt := range tests {
		succeeded := checkAltNamesExist(rt.IPs, rt.DNSNames, rt.requiredAltNames)
		if succeeded != rt.succeed {
			t.Errorf(
				"failed checkAltNamesExist:\n\texpected: %t\n\t  actual: %t",
				rt.succeed,
				succeeded,
			)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017-21 the original author or 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 pbf // import "m4o.io/pbf"

import "strconv"

// ftoa converts a float to a string with no trailing zeros.
// Snagged from github.com/dustin/go-humanize to minimize library dependencies.
func ftoa(num float64) string {
	return stripTrailingZeros(strconv.FormatFloat(num, 'f', 6, 64))
}

func stripTrailingZeros(s string) string {
	offset := len(s) - 1
	for offset > 0 {
		if s[offset] == '.' {
			offset--
			break
		}
		if s[offset] != '0' {
			break
		}
		offset--
	}
	return s[:offset+1]
}

<commit_msg>Lint cleanup of fmt.go<commit_after>// Copyright 2017-21 the original author or 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 pbf // import "m4o.io/pbf"

import "strconv"

// ftoa converts a float to a string with no trailing zeros.
// Snagged from github.com/dustin/go-humanize to minimize library dependencies.
func ftoa(num float64) string {
	return stripTrailingZeros(strconv.FormatFloat(num, 'f', 6, 64))
}

func stripTrailingZeros(s string) string {
	offset := len(s) - 1
	for offset > 0 {
		if s[offset] == '.' {
			offset--
			break
		}

		if s[offset] != '0' {
			break
		}

		offset--
	}

	return s[:offset+1]
}
<|endoftext|>


<commit_before>package typewriter

import (
	"bytes"
	"fmt"
)

const (
	colorOff     = "\033[0m"
	colorDefault = "\033[41m"
)

type Config struct {
	ShowLineNumbers bool

	// If true, marks the first spot where the two columns diverge.
	MarkFirstDifference bool

	// The escape sequence to use to mark a difference. If not specified
	// defaults to a red background.
	Marking string

	Padding   int
	Separator string

	LeftHeader  string
	RightHeader string
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Returns the length of the longest string.
func maxWidth(lines []string) int {
	max := 0
	for _, line := range lines {
		max = maxInt(len(line), max)
	}
	return max
}

func getOrEmpty(strings []string, index int) string {
	if index >= len(strings) {
		return ""
	}
	return strings[index]
}

func rightZero(s string, desiredLength int, padChar string) string {
	if len(s) > desiredLength {
		panic(fmt.Sprintf("Length of string (%d) is greater than desired length (%d)", len(s), desiredLength))
	}

	// TODO: Optimize
	padding := desiredLength - len(s)
	for i := 0; i < padding; i++ {
		s = s + padChar
	}

	return s
}

func leftZero(s string, desiredLength int, padChar string) string {
	if len(s) > desiredLength {
		panic(fmt.Sprintf("Length of string (%d) is greater than desired length (%d)", len(s), desiredLength))
	}

	// TODO: Optimize
	padding := desiredLength - len(s)
	for i := 0; i < padding; i++ {
		s = padChar + s
	}

	return s
}

func findDifference(s1, s2 string) (index int, found bool) {
	minLen := minInt(len(s1), len(s2))
	for i := 0; i < minLen; i++ {
		c1 := s1[i]
		c2 := s2[i]

		if c1 != c2 {
			return i, true
		}
	}

	if len(s1) > len(s2) {
		return len(s2), true
	} else if len(s2) > len(s1) {
		return len(s1), true
	}

	return -1, false
}

func Run(lines1, lines2 []string, config Config) string {
	var buf bytes.Buffer

	leftColumnWidth := maxWidth(append(lines1, config.LeftHeader))

	padding := rightZero("", config.Padding, " ")

	if config.Marking == "" {
		config.Marking = colorDefault
	}

	// Find the max width we'll need for line numbers.
	maxLineNumber := maxInt(len(lines1), len(lines2))
	maxLineNumberWidth := len(fmt.Sprintf("%d", maxLineNumber)) + 2

	if config.LeftHeader != "" || config.RightHeader != "" {
		h1 := rightZero(config.LeftHeader, leftColumnWidth, " ")
		lineNumber := ""
		if config.ShowLineNumbers {
			lineNumber = leftZero(lineNumber, maxLineNumberWidth, " ")
		}
		buf.WriteString(lineNumber + h1 + padding + config.Separator + config.RightHeader + "\n\n")
	}

	differenceFound := false
	for i := 0; i < maxInt(len(lines1), len(lines2)); i++ {
		s1 := getOrEmpty(lines1, i)
		s2 := getOrEmpty(lines2, i)

		extraWidth := 0
		if config.MarkFirstDifference && !differenceFound {
			diffIndex, found := findDifference(s1, s2)
			if found {
				differenceFound = true
				s1 = s1[:diffIndex] + config.Marking + s1[diffIndex:] + colorOff
				s2 = s2[:diffIndex] + config.Marking + s2[diffIndex:] + colorOff
				extraWidth = len(config.Marking) + len(colorOff)
			}
		}

		s1 = rightZero(s1, leftColumnWidth+extraWidth, " ")

		lineNumber := ""
		if config.ShowLineNumbers {
			lineNumber = fmt.Sprint(i+1) + ". "
			lineNumber = leftZero(lineNumber, maxLineNumberWidth, " ")
		}

		buf.WriteString(lineNumber + s1 + padding + config.Separator + s2 + "\n")
	}

	return buf.String()
}
<commit_msg>Speed up converting line numbers to strings<commit_after>package typewriter

import (
	"bytes"
	"fmt"
	"strconv"
)

const (
	colorOff     = "\033[0m"
	colorDefault = "\033[41m"
)

type Config struct {
	ShowLineNumbers bool

	// If true, marks the first spot where the two columns diverge.
	MarkFirstDifference bool

	// The escape sequence to use to mark a difference. If not specified
	// defaults to a red background.
	Marking string

	Padding   int
	Separator string

	LeftHeader  string
	RightHeader string
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// Returns the length of the longest string.
func maxWidth(lines []string) int {
	max := 0
	for _, line := range lines {
		max = maxInt(len(line), max)
	}
	return max
}

func getOrEmpty(strings []string, index int) string {
	if index >= len(strings) {
		return ""
	}
	return strings[index]
}

func rightZero(s string, desiredLength int, padChar string) string {
	if len(s) > desiredLength {
		panic(fmt.Sprintf("Length of string (%d) is greater than desired length (%d)", len(s), desiredLength))
	}

	// TODO: Optimize
	padding := desiredLength - len(s)
	for i := 0; i < padding; i++ {
		s = s + padChar
	}

	return s
}

func leftZero(s string, desiredLength int, padChar string) string {
	if len(s) > desiredLength {
		panic(fmt.Sprintf("Length of string (%d) is greater than desired length (%d)", len(s), desiredLength))
	}

	// TODO: Optimize
	padding := desiredLength - len(s)
	for i := 0; i < padding; i++ {
		s = padChar + s
	}

	return s
}

func findDifference(s1, s2 string) (index int, found bool) {
	minLen := minInt(len(s1), len(s2))
	for i := 0; i < minLen; i++ {
		c1 := s1[i]
		c2 := s2[i]

		if c1 != c2 {
			return i, true
		}
	}

	if len(s1) > len(s2) {
		return len(s2), true
	} else if len(s2) > len(s1) {
		return len(s1), true
	}

	return -1, false
}

func Run(lines1, lines2 []string, config Config) string {
	var buf bytes.Buffer

	leftColumnWidth := maxWidth(append(lines1, config.LeftHeader))

	padding := rightZero("", config.Padding, " ")

	if config.Marking == "" {
		config.Marking = colorDefault
	}

	// Find the max width we'll need for line numbers.
	maxLineNumber := maxInt(len(lines1), len(lines2))
	maxLineNumberWidth := len(fmt.Sprintf("%d", maxLineNumber)) + 2

	if config.LeftHeader != "" || config.RightHeader != "" {
		h1 := rightZero(config.LeftHeader, leftColumnWidth, " ")
		lineNumber := ""
		if config.ShowLineNumbers {
			lineNumber = leftZero(lineNumber, maxLineNumberWidth, " ")
		}
		buf.WriteString(lineNumber + h1 + padding + config.Separator + config.RightHeader + "\n\n")
	}

	differenceFound := false
	for i := 0; i < maxInt(len(lines1), len(lines2)); i++ {
		s1 := getOrEmpty(lines1, i)
		s2 := getOrEmpty(lines2, i)

		extraWidth := 0
		if config.MarkFirstDifference && !differenceFound {
			diffIndex, found := findDifference(s1, s2)
			if found {
				differenceFound = true
				s1 = s1[:diffIndex] + config.Marking + s1[diffIndex:] + colorOff
				s2 = s2[:diffIndex] + config.Marking + s2[diffIndex:] + colorOff
				extraWidth = len(config.Marking) + len(colorOff)
			}
		}

		s1 = rightZero(s1, leftColumnWidth+extraWidth, " ")

		lineNumber := ""
		if config.ShowLineNumbers {
			lineNumber = strconv.Itoa(i + 1) + ". "
			lineNumber = leftZero(lineNumber, maxLineNumberWidth, " ")
		}

		buf.WriteString(lineNumber + s1 + padding + config.Separator + s2 + "\n")
	}

	return buf.String()
}
<|endoftext|>


<commit_before>package goftp

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
)

var REGEX_PWD_PATH *regexp.Regexp = regexp.MustCompile(`\"(.*)\"`)

type FTP struct {
	conn net.Conn

	addr string

	debug     bool
	tlsconfig *tls.Config

	reader *bufio.Reader
	writer *bufio.Writer
}

func (ftp *FTP) Close() {
	ftp.conn.Close()
}

type WalkFunc func(path string, info os.FileMode, err error) error
type RetrFunc func(r io.Reader) error

func parseLine(line string) (perm string, t string, filename string) {
	for _, v := range strings.Split(line, ";") {
		v2 := strings.Split(v, "=")

		switch v2[0] {
		case "perm":
			perm = v2[1]
		case "type":
			t = v2[1]
		default:
			filename = v[1 : len(v)-2]
		}
	}

	return
}

// walks recursively through path and call walkfunc for each file
func (ftp *FTP) Walk(path string, walkFn WalkFunc) (err error) {
	/*
		if err = walkFn(path, os.ModeDir, nil); err != nil {
			if err == filepath.SkipDir {
				return nil
			}
		}
	*/
	if ftp.debug {
		log.Printf("Walking: '%s'\n", path)
	}

	var lines []string

	if lines, err = ftp.List(path); err != nil {
		return
	}

	for _, line := range lines {
		_, t, subpath := parseLine(line)

		switch t {
		case "dir":
			if subpath == "." {
			} else if subpath == ".." {
			} else {
				if err = ftp.Walk(path+subpath+"/", walkFn); err != nil {
					return
				}
			}
		case "file":
			if err = walkFn(path+subpath, os.FileMode(0), nil); err != nil {
				return
			}
		}
	}

	return
}

// send quit to the server and close the connection
func (ftp *FTP) Quit() (err error) {
	if _, err := ftp.cmd("221", "QUIT"); err != nil {
		return err
	}

	ftp.conn.Close()
	ftp.conn = nil

	return nil
}

// will send a NOOP (no operation) to the server
func (ftp *FTP) Noop() (err error) {
	_, err = ftp.cmd("200", "NOOP")
	return
}

// private function to send command and compare return code with expects
func (ftp *FTP) cmd(expects string, command string, args ...interface{}) (line string, err error) {
	if err = ftp.send(command, args...); err != nil {
		return
	}

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, expects) {
		err = errors.New(line)
		return
	}

	return
}

// rename file
func (ftp *FTP) Rename(from string, to string) (err error) {
	if _, err = ftp.cmd("350", "RNFR %s", from); err != nil {
		return
	}

	if _, err = ftp.cmd("250", "RNTO %s", to); err != nil {
		return
	}

	return
}

// make directory
func (ftp *FTP) Mkd(path string) error {
	_, err := ftp.cmd("257", "MKD %s", path)
	return err
}

// get current path
func (ftp *FTP) Pwd() (path string, err error) {
	var line string
	if line, err = ftp.cmd("257", "PWD"); err != nil {
		return
	}

	res := REGEX_PWD_PATH.FindAllStringSubmatch(line[4:], -1)

	path = res[0][1]
	return
}

// change current path
func (ftp *FTP) Cwd(path string) (err error) {
	_, err = ftp.cmd("250", "CWD %s", path)
	return
}

// delete file
func (ftp *FTP) Dele(path string) (err error) {
	if err = ftp.send("DELE %s", path); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "250") {
		return errors.New(line)
	}

	return
}

// secures the ftp connection by using TLS
func (ftp *FTP) AuthTLS(config tls.Config) error {
	if _, err := ftp.cmd("234", "AUTH TLS"); err != nil {
		return err
	}

	// wrap tls on existing connection
	ftp.tlsconfig = &config

	ftp.conn = tls.Client(ftp.conn, &config)
	ftp.writer = bufio.NewWriter(ftp.conn)
	ftp.reader = bufio.NewReader(ftp.conn)

	if _, err := ftp.cmd("200", "PROT P"); err != nil {
		return err
	}

	return nil
}

// change transfer type
func (ftp *FTP) Type(t string) error {
	_, err := ftp.cmd("200", "TYPE %s", t)
	return err
}

func (ftp *FTP) receive() (string, error) {
	line, err := ftp.reader.ReadString('\n')

	if ftp.debug {
		log.Printf("< %s", line)
	}

	return line, err
}

func (ftp *FTP) send(command string, arguments ...interface{}) error {
	if ftp.debug {
		log.Printf("> %s", fmt.Sprintf(command, arguments...))
	}

	command = fmt.Sprintf(command, arguments...)
	command += "\r\n"

	if _, err := ftp.writer.WriteString(command); err != nil {
		return err
	}

	if err := ftp.writer.Flush(); err != nil {
		return err
	}

	return nil
}

// enables passive data connection and returns port number
func (ftp *FTP) Pasv() (port int, err error) {
	var line string
	if line, err = ftp.cmd("227", "PASV"); err != nil {
		return
	}

	re, err := regexp.Compile(`\((.*)\)`)

	res := re.FindAllStringSubmatch(line, -1)

	s := strings.Split(res[0][1], ",")

	l1, _ := strconv.Atoi(s[len(s)-2])
	l2, _ := strconv.Atoi(s[len(s)-1])

	port = l1<<8 + l2

	return
}

// open new data connection
func (ftp *FTP) newConnection(port int) (conn net.Conn, err error) {
	addr := fmt.Sprintf("%s:%d", strings.Split(ftp.addr, ":")[0], port)

	if ftp.debug {
		log.Printf("Connecting to %s\n", addr)
	}

	if conn, err = net.Dial("tcp", addr); err != nil {
		return
	}

	if ftp.tlsconfig != nil {
		conn = tls.Client(conn, ftp.tlsconfig)
	}

	return
}

// upload file
func (ftp *FTP) Stor(path string, r io.Reader) (err error) {
	if err = ftp.Type("I"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	if err = ftp.send("STOR %s", path); err != nil {
		return
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	if _, err = io.Copy(pconn, r); err != nil {
		return
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return

}

// retrieves file
func (ftp *FTP) Retr(path string, retrFn RetrFunc) (s string, err error) {
	if err = ftp.Type("I"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	if err = ftp.send("RETR %s", path); err != nil {
		return
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	if err = retrFn(pconn); err != nil {
		return
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return
}

// list the path (or current directory)
func (ftp *FTP) List(path string) (files []string, err error) {
	if err = ftp.Type("A"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	// _, err = ftp.writer.WriteString(fmt.Sprintf("LIST %s\r\n", path))
	// check for features LIST / MLSD
	if err = ftp.send("MLSD %s", path); err != nil {
		return
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	reader := bufio.NewReader(pconn)

	for {
		line, err = reader.ReadString('\n')

		if err == io.EOF {
			break
		} else if err != nil {
			return
		}

		files = append(files, string(line))
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return
}

// login to the server
func (ftp *FTP) Login(username string, password string) (err error) {
	if _, err = ftp.cmd("331", "USER %s", username); err != nil {
		return
	}

	if _, err = ftp.cmd("230", "PASS %s", password); err != nil {
		return
	}

	return
}

// connect to server
func Connect(addr string) (*FTP, error) {
	var err error
	var conn net.Conn

	if conn, err = net.Dial("tcp", addr); err != nil {
		return nil, err
	}

	writer := bufio.NewWriter(conn)
	reader := bufio.NewReader(conn)

	var line string

	line, err = reader.ReadString('\n')

	var debug bool = true

	if debug {
		log.Print(line)
	}

	return &FTP{conn: conn, addr: addr, reader: reader, writer: writer, debug: debug}, nil
}
<commit_msg>Modified for anonymous and debug<commit_after>package goftp

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
)

var REGEX_PWD_PATH *regexp.Regexp = regexp.MustCompile(`\"(.*)\"`)

type FTP struct {
	conn net.Conn

	addr string

	debug     bool
	tlsconfig *tls.Config

	reader *bufio.Reader
	writer *bufio.Writer
}

func (ftp *FTP) Close() {
	ftp.conn.Close()
}

type WalkFunc func(path string, info os.FileMode, err error) error
type RetrFunc func(r io.Reader) error

func parseLine(line string) (perm string, t string, filename string) {
	for _, v := range strings.Split(line, ";") {
		v2 := strings.Split(v, "=")

		switch v2[0] {
		case "perm":
			perm = v2[1]
		case "type":
			t = v2[1]
		default:
			filename = v[1 : len(v)-2]
		}
	}

	return
}

// walks recursively through path and call walkfunc for each file
func (ftp *FTP) Walk(path string, walkFn WalkFunc) (err error) {
	/*
		if err = walkFn(path, os.ModeDir, nil); err != nil {
			if err == filepath.SkipDir {
				return nil
			}
		}
	*/
	if ftp.debug {
		log.Printf("Walking: '%s'\n", path)
	}

	var lines []string

	if lines, err = ftp.List(path); err != nil {
		return
	}

	for _, line := range lines {
		_, t, subpath := parseLine(line)

		switch t {
		case "dir":
			if subpath == "." {
			} else if subpath == ".." {
			} else {
				if err = ftp.Walk(path+subpath+"/", walkFn); err != nil {
					return
				}
			}
		case "file":
			if err = walkFn(path+subpath, os.FileMode(0), nil); err != nil {
				return
			}
		}
	}

	return
}

// send quit to the server and close the connection
func (ftp *FTP) Quit() (err error) {
	if _, err := ftp.cmd("221", "QUIT"); err != nil {
		return err
	}

	ftp.conn.Close()
	ftp.conn = nil

	return nil
}

// will send a NOOP (no operation) to the server
func (ftp *FTP) Noop() (err error) {
	_, err = ftp.cmd("200", "NOOP")
	return
}

// private function to send command and compare return code with expects
func (ftp *FTP) cmd(expects string, command string, args ...interface{}) (line string, err error) {
	if err = ftp.send(command, args...); err != nil {
		return
	}

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, expects) {
		err = errors.New(line)
		return
	}

	return
}

// rename file
func (ftp *FTP) Rename(from string, to string) (err error) {
	if _, err = ftp.cmd("350", "RNFR %s", from); err != nil {
		return
	}

	if _, err = ftp.cmd("250", "RNTO %s", to); err != nil {
		return
	}

	return
}

// make directory
func (ftp *FTP) Mkd(path string) error {
	_, err := ftp.cmd("257", "MKD %s", path)
	return err
}

// get current path
func (ftp *FTP) Pwd() (path string, err error) {
	var line string
	if line, err = ftp.cmd("257", "PWD"); err != nil {
		return
	}

	res := REGEX_PWD_PATH.FindAllStringSubmatch(line[4:], -1)

	path = res[0][1]
	return
}

// change current path
func (ftp *FTP) Cwd(path string) (err error) {
	_, err = ftp.cmd("250", "CWD %s", path)
	return
}

// delete file
func (ftp *FTP) Dele(path string) (err error) {
	if err = ftp.send("DELE %s", path); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "250") {
		return errors.New(line)
	}

	return
}

// secures the ftp connection by using TLS
func (ftp *FTP) AuthTLS(config tls.Config) error {
	if _, err := ftp.cmd("234", "AUTH TLS"); err != nil {
		return err
	}

	// wrap tls on existing connection
	ftp.tlsconfig = &config

	ftp.conn = tls.Client(ftp.conn, &config)
	ftp.writer = bufio.NewWriter(ftp.conn)
	ftp.reader = bufio.NewReader(ftp.conn)

	if _, err := ftp.cmd("200", "PROT P"); err != nil {
		return err
	}

	return nil
}

// change transfer type
func (ftp *FTP) Type(t string) error {
	_, err := ftp.cmd("200", "TYPE %s", t)
	return err
}

func (ftp *FTP) receive() (string, error) {
	line, err := ftp.reader.ReadString('\n')

	if ftp.debug {
		log.Printf("< %s", line)
	}

	return line, err
}

func (ftp *FTP) send(command string, arguments ...interface{}) error {
	if ftp.debug {
		log.Printf("> %s", fmt.Sprintf(command, arguments...))
	}

	command = fmt.Sprintf(command, arguments...)
	command += "\r\n"

	if _, err := ftp.writer.WriteString(command); err != nil {
		return err
	}

	if err := ftp.writer.Flush(); err != nil {
		return err
	}

	return nil
}

// enables passive data connection and returns port number
func (ftp *FTP) Pasv() (port int, err error) {
	var line string
	if line, err = ftp.cmd("227", "PASV"); err != nil {
		return
	}

	re, err := regexp.Compile(`\((.*)\)`)

	res := re.FindAllStringSubmatch(line, -1)

	s := strings.Split(res[0][1], ",")

	l1, _ := strconv.Atoi(s[len(s)-2])
	l2, _ := strconv.Atoi(s[len(s)-1])

	port = l1<<8 + l2

	return
}

// open new data connection
func (ftp *FTP) newConnection(port int) (conn net.Conn, err error) {
	addr := fmt.Sprintf("%s:%d", strings.Split(ftp.addr, ":")[0], port)

	if ftp.debug {
		log.Printf("Connecting to %s\n", addr)
	}

	if conn, err = net.Dial("tcp", addr); err != nil {
		return
	}

	if ftp.tlsconfig != nil {
		conn = tls.Client(conn, ftp.tlsconfig)
	}

	return
}

// upload file
func (ftp *FTP) Stor(path string, r io.Reader) (err error) {
	if err = ftp.Type("I"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	if err = ftp.send("STOR %s", path); err != nil {
		return
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	if _, err = io.Copy(pconn, r); err != nil {
		return
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return

}

// retrieves file
func (ftp *FTP) Retr(path string, retrFn RetrFunc) (s string, err error) {
	if err = ftp.Type("I"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	if err = ftp.send("RETR %s", path); err != nil {
		return
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	if err = retrFn(pconn); err != nil {
		return
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return
}

// list the path (or current directory)
func (ftp *FTP) List(path string) (files []string, err error) {
	if err = ftp.Type("A"); err != nil {
		return
	}

	var port int
	if port, err = ftp.Pasv(); err != nil {
		return
	}

	// check if MLSD works
	if err = ftp.send("MLSD %s", path); err != nil {
		if err = ftp.send("LIST %s", path); err != nil {
			return
		}
	}

	var pconn net.Conn
	if pconn, err = ftp.newConnection(port); err != nil {
		return
	}

	var line string
	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "150") {
		err = errors.New(line)
		return
	}

	reader := bufio.NewReader(pconn)

	for {
		line, err = reader.ReadString('\n')

		if err == io.EOF {
			break
		} else if err != nil {
			return
		}

		files = append(files, string(line))
	}

	pconn.Close()

	if line, err = ftp.receive(); err != nil {
		return
	}

	if !strings.HasPrefix(line, "226") {
		err = errors.New(line)
		return
	}

	return
}

// login to the server
func (ftp *FTP) Login(username string, password string) (err error) {
	if _, err = ftp.cmd("331", "USER %s", username); err != nil {
		if strings.HasPrefix(err.Error(), "230") {
			// Ok, probably anonymous server
			// but login was fine, so return no error
			err = nil
		} else {
			return
		}
	}

	if _, err = ftp.cmd("230", "PASS %s", password); err != nil {
		return
	}

	return
}

// connect to server
func Connect(addr string, debugp bool) (*FTP, error) {
	var err error
	var conn net.Conn

	if conn, err = net.Dial("tcp", addr); err != nil {
		return nil, err
	}

	writer := bufio.NewWriter(conn)
	reader := bufio.NewReader(conn)

	var line string

	line, err = reader.ReadString('\n')

	var debug bool = debugp

	if debug {
		log.Print(line)
	}

	return &FTP{conn: conn, addr: addr, reader: reader, writer: writer, debug: debug}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net"
	"net/url"
	"strconv"
	"time"
)

const (
	connectMinResponseLength = 16
	announceMinResponseLength = 20
	connectBufferSize = 150
	announceBufferSize = 20000
	startPort = 6881
	endPort = 6889
)

type UdpTracker struct {
	*tracker
	connectionId  uint64
	transactionId uint32
	serverAddr    *net.UDPAddr
	conn          *net.UDPConn
}

type connectRequest struct {
	connectionId  uint64
	action        uint32
	transactionId uint32
}

type connectResponse struct {
	action        uint32
	transactionId uint32
	connectionId  uint64
}

type errorResponse struct {
	action        uint32
	transactionId uint32
	message       string
}

type announceRequest struct {
	connectionId  uint64
	action        uint32
	transactionId uint32
	infoHash      []byte
	peerId        []byte
	downloaded    uint64
	left          uint64
	uploaded      uint64
	event         uint32
	ipAddr        uint32
	key           uint32
	numWant       int32
	port          uint16
}

type announceResponse struct {
	action        uint32
	transactionId uint32
	interval      uint32
	leechers      uint32
	seeders       uint32
	peers         []PeerTuple
}

func (r *connectRequest) MarshalBinary() ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, r.connectionId)
	err = binary.Write(buf, binary.BigEndian, r.action)
	err = binary.Write(buf, binary.BigEndian, r.transactionId)
	return buf.Bytes(), err
}

func (r *connectResponse) UnmarshalBinary(data []byte) error {
	buf := bytes.NewReader(data[4:])
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(buf, binary.BigEndian, &r.transactionId)
	err = binary.Read(buf, binary.BigEndian, &r.connectionId)
	return err
}

func (r *connectResponse) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprintln(&b, r.connectionId, r.action, r.transactionId)
	return b.Bytes(), nil
}

func (r *errorResponse) UnmarshalBinary(data []byte) error {
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(bytes.NewReader(data[4:]), binary.BigEndian, &r.transactionId)
	r.message = string(data[8:])
	return err
}

func (r *announceRequest) MarshalBinary() ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, r.connectionId)
	err = binary.Write(buf, binary.BigEndian, r.action)
	err = binary.Write(buf, binary.BigEndian, r.transactionId)
	err = binary.Write(buf, binary.BigEndian, r.infoHash)
	err = binary.Write(buf, binary.BigEndian, r.peerId)
	err = binary.Write(buf, binary.BigEndian, r.downloaded)
	err = binary.Write(buf, binary.BigEndian, r.left)
	err = binary.Write(buf, binary.BigEndian, r.uploaded)
	err = binary.Write(buf, binary.BigEndian, r.event)
	err = binary.Write(buf, binary.BigEndian, r.ipAddr)
	err = binary.Write(buf, binary.BigEndian, r.key)
	err = binary.Write(buf, binary.BigEndian, r.numWant)
	err = binary.Write(buf, binary.BigEndian, r.port)
	return buf.Bytes(), err
}

func (r *announceResponse) UnmarshalBinary(data []byte) error {
	buf := bytes.NewReader(data[4:])
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(buf, binary.BigEndian, &r.transactionId)
	err = binary.Read(buf, binary.BigEndian, &r.interval)
	err = binary.Read(buf, binary.BigEndian, &r.leechers)
	err = binary.Read(buf, binary.BigEndian, &r.seeders)

	if len(data) > announceMinResponseLength {
		peerBytes := data[announceMinResponseLength:]
		peers := len(peerBytes) / 6

		ipLen := 4
		portLen := 2
		for i := 0; i < peers; i++ {
			ipOffset := i*6
			portOffset := ipOffset + ipLen
			peer := PeerTuple{}
			ipBytes := peerBytes[ipOffset : ipOffset + ipLen]
			peer.IP = net.IPv4(ipBytes[0], ipBytes[1], ipBytes[2], ipBytes[3])
			binary.Read(bytes.NewReader(peerBytes[portOffset : portOffset + portLen]), binary.BigEndian, &peer.Port)
			r.peers = append(r.peers, peer)
		}
	}

	return err
}

func NewUdpTracker(key string, chans trackerPeerChans, port uint16, infoHash []byte, announce *url.URL) *UdpTracker {
	return &UdpTracker{&tracker{key: key, peerChans: chans, port: port, infoHash: infoHash, announceURL: announce}, 0, 0, &net.UDPAddr{}, &net.UDPConn{}}
}

func (tr *UdpTracker) Announce(event int) {
	err := tr.connect()
	if err != nil {
		log.Printf("Tracker : Could not connect to tracker %s: \n", tr.announceURL.String(), err)
		return
	}

	key, _ := strconv.ParseUint(tr.key, 16, 4)
	announce := &announceRequest{
		connectionId:  tr.connectionId,
		action:        1,
		transactionId: tr.transactionId,
		infoHash:      tr.infoHash,
		peerId:        PeerID[:],
		downloaded:    uint64(tr.stats.Downloaded),
		left:          uint64(tr.stats.Left),
		uploaded:      uint64(tr.stats.Uploaded),
		event:         uint32(event),
		ipAddr:        0,
		key:           uint32(key),
		numWant:       -1,
		port:          6881,
	}
	announceBytes, _ := announce.MarshalBinary()

	buff := make([]byte, announceBufferSize)
	length := tr.request(announceBytes, buff)

	if length >= announceMinResponseLength {
		var response announceResponse
		response.UnmarshalBinary(buff[:length])

		if event != Stopped {
			if response.interval != 0 {
				nextAnnounce := time.Second * time.Duration(response.interval)
				log.Println("Tracker : Announce : Scheduling next announce in", nextAnnounce)
				tr.timer = time.After(nextAnnounce)
			}

			for _, peer := range response.peers {
				tr.peerChans.peers <- peer
			}
		}
	}
}

func (tr *UdpTracker) connect() error {
	log.Printf("Tracker : Connect (%v)", tr.announceURL)

	connectReq := connectRequest{connectionId: 0x41727101980, action: 0, transactionId: tr.transactionId}
	connectBytes, _ := connectReq.MarshalBinary()

	buff := make([]byte, connectBufferSize)
	length := tr.request(connectBytes, buff)

	if length >= connectMinResponseLength {
		var response connectResponse
		response.UnmarshalBinary(buff[:length])

		if response.action == 3 || tr.transactionId != response.transactionId {
			error := errorResponse{}
			error.UnmarshalBinary(buff)
			return errors.New("Response Error: " + error.message)
		}

		tr.connectionId = response.connectionId
		return nil
	}

	return errors.New("Invalid connect response length")
}

func (tr *UdpTracker) Run() {
	log.Printf("Tracker : Run : Started (%s)\n", tr.announceURL)
	defer log.Printf("Tracker : Run : Completed (%s)\n", tr.announceURL)

	rand.Seed(time.Now().UnixNano())

	serverAddr, err := net.ResolveUDPAddr("udp", tr.announceURL.Host)
	if err != nil {
		log.Println("Tracker : Could not resolve tracker host!")
		return
	}

	var conn *net.UDPConn
	for port := startPort; err == nil && port < endPort; port++ {
		conn, err = net.ListenUDP("udp", &net.UDPAddr{Port: port})
	}
	if err != nil {
		log.Printf("Tracker : Could not bind to any port in range %d-%d\n", startPort, endPort)
		return
	}

	tr.transactionId = rand.Uint32()
	tr.serverAddr = serverAddr
	tr.conn = conn
	tr.Announce(Started)

	for {
		select {
		case <-tr.quit:
			log.Println("Tracker : Stop : Stopping")
			tr.Announce(Stopped)
			return
		case <-tr.completedCh:
			go tr.Announce(Completed)
		case <-tr.timer:
			log.Printf("Tracker : Run : Interval Timer Expired (%s)\n", tr.announceURL)
			go tr.Announce(Interval)
		case stats := <-tr.peerChans.stats:
			log.Println("read from stats", stats)
		}
	}
}

// Send a udp packet to the tracker, fill in the dest buffer with the response
func (tr *UdpTracker) request(payload []byte, dest []byte) int {
	n := 0
	totalAttempts := 0

	// notify the sender when a response is recieved
	recvChan := make(chan bool)

	// keep sending the packet and wait for the specifed time
	// before trying again, limit n to 8
	go func() {

		// initial send
		tr.conn.WriteTo(payload, tr.serverAddr)
	Listen:
		for {
			// timeout: 15 * 2 ^ n (0-8)
			timeout := time.Second * time.Duration(15 * int(math.Pow(2.0, float64(n))))
			timer := time.After(timeout)
			select {
			case <-recvChan:
				break Listen
			case <-timer:
				tr.conn.WriteTo(payload, tr.serverAddr)
				totalAttempts++
				if n < 8 { n++ }
			}
		}
	}()

	// read until some data is received
	var err error
	length := 0
	for length == 0 {
		length, _, err = tr.conn.ReadFrom(dest)

		// not sure what to do here?
		if err != nil {
			log.Println("Tracker : udp packet read error")
		}
	}

	recvChan <- true
	return length
}
<commit_msg>Removed arbitrary port range<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net"
	"net/url"
	"strconv"
	"time"
)

const (
	connectMinResponseLength = 16
	announceMinResponseLength = 20
	connectBufferSize = 150
	announceBufferSize = 20000
)

type UdpTracker struct {
	*tracker
	connectionId  uint64
	transactionId uint32
	serverAddr    *net.UDPAddr
	conn          *net.UDPConn
}

type connectRequest struct {
	connectionId  uint64
	action        uint32
	transactionId uint32
}

type connectResponse struct {
	action        uint32
	transactionId uint32
	connectionId  uint64
}

type errorResponse struct {
	action        uint32
	transactionId uint32
	message       string
}

type announceRequest struct {
	connectionId  uint64
	action        uint32
	transactionId uint32
	infoHash      []byte
	peerId        []byte
	downloaded    uint64
	left          uint64
	uploaded      uint64
	event         uint32
	ipAddr        uint32
	key           uint32
	numWant       int32
	port          uint16
}

type announceResponse struct {
	action        uint32
	transactionId uint32
	interval      uint32
	leechers      uint32
	seeders       uint32
	peers         []PeerTuple
}

func (r *connectRequest) MarshalBinary() ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, r.connectionId)
	err = binary.Write(buf, binary.BigEndian, r.action)
	err = binary.Write(buf, binary.BigEndian, r.transactionId)
	return buf.Bytes(), err
}

func (r *connectResponse) UnmarshalBinary(data []byte) error {
	buf := bytes.NewReader(data[4:])
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(buf, binary.BigEndian, &r.transactionId)
	err = binary.Read(buf, binary.BigEndian, &r.connectionId)
	return err
}

func (r *connectResponse) MarshalBinary() ([]byte, error) {
	var b bytes.Buffer
	fmt.Fprintln(&b, r.connectionId, r.action, r.transactionId)
	return b.Bytes(), nil
}

func (r *errorResponse) UnmarshalBinary(data []byte) error {
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(bytes.NewReader(data[4:]), binary.BigEndian, &r.transactionId)
	r.message = string(data[8:])
	return err
}

func (r *announceRequest) MarshalBinary() ([]byte, error) {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, r.connectionId)
	err = binary.Write(buf, binary.BigEndian, r.action)
	err = binary.Write(buf, binary.BigEndian, r.transactionId)
	err = binary.Write(buf, binary.BigEndian, r.infoHash)
	err = binary.Write(buf, binary.BigEndian, r.peerId)
	err = binary.Write(buf, binary.BigEndian, r.downloaded)
	err = binary.Write(buf, binary.BigEndian, r.left)
	err = binary.Write(buf, binary.BigEndian, r.uploaded)
	err = binary.Write(buf, binary.BigEndian, r.event)
	err = binary.Write(buf, binary.BigEndian, r.ipAddr)
	err = binary.Write(buf, binary.BigEndian, r.key)
	err = binary.Write(buf, binary.BigEndian, r.numWant)
	err = binary.Write(buf, binary.BigEndian, r.port)
	return buf.Bytes(), err
}

func (r *announceResponse) UnmarshalBinary(data []byte) error {
	buf := bytes.NewReader(data[4:])
	actionBytes := append([]byte{0, 0, 0}, data[0])
	err := binary.Read(bytes.NewReader(actionBytes), binary.BigEndian, &r.action)
	err = binary.Read(buf, binary.BigEndian, &r.transactionId)
	err = binary.Read(buf, binary.BigEndian, &r.interval)
	err = binary.Read(buf, binary.BigEndian, &r.leechers)
	err = binary.Read(buf, binary.BigEndian, &r.seeders)

	if len(data) > announceMinResponseLength {
		peerBytes := data[announceMinResponseLength:]
		peers := len(peerBytes) / 6

		ipLen := 4
		portLen := 2
		for i := 0; i < peers; i++ {
			ipOffset := i*6
			portOffset := ipOffset + ipLen
			peer := PeerTuple{}
			ipBytes := peerBytes[ipOffset : ipOffset + ipLen]
			peer.IP = net.IPv4(ipBytes[0], ipBytes[1], ipBytes[2], ipBytes[3])
			binary.Read(bytes.NewReader(peerBytes[portOffset : portOffset + portLen]), binary.BigEndian, &peer.Port)
			r.peers = append(r.peers, peer)
		}
	}

	return err
}

func NewUdpTracker(key string, chans trackerPeerChans, port uint16, infoHash []byte, announce *url.URL) *UdpTracker {
	return &UdpTracker{&tracker{key: key, peerChans: chans, port: port, infoHash: infoHash, announceURL: announce}, 0, 0, &net.UDPAddr{}, &net.UDPConn{}}
}

func (tr *UdpTracker) Announce(event int) {
	err := tr.connect()
	if err != nil {
		log.Printf("Tracker : Could not connect to tracker %s: \n", tr.announceURL.String(), err)
		return
	}

	key, _ := strconv.ParseUint(tr.key, 16, 4)
	announce := &announceRequest{
		connectionId:  tr.connectionId,
		action:        1,
		transactionId: tr.transactionId,
		infoHash:      tr.infoHash,
		peerId:        PeerID[:],
		downloaded:    uint64(tr.stats.Downloaded),
		left:          uint64(tr.stats.Left),
		uploaded:      uint64(tr.stats.Uploaded),
		event:         uint32(event),
		ipAddr:        0,
		key:           uint32(key),
		numWant:       -1,
		port:          6881,
	}
	announceBytes, _ := announce.MarshalBinary()

	buff := make([]byte, announceBufferSize)
	length := tr.request(announceBytes, buff)

	if length >= announceMinResponseLength {
		var response announceResponse
		response.UnmarshalBinary(buff[:length])

		if event != Stopped {
			if response.interval != 0 {
				nextAnnounce := time.Second * time.Duration(response.interval)
				log.Println("Tracker : Announce : Scheduling next announce in", nextAnnounce)
				tr.timer = time.After(nextAnnounce)
			}

			for _, peer := range response.peers {
				tr.peerChans.peers <- peer
			}
		}
	}
}

func (tr *UdpTracker) connect() error {
	log.Printf("Tracker : Connect (%v)", tr.announceURL)

	connectReq := connectRequest{connectionId: 0x41727101980, action: 0, transactionId: tr.transactionId}
	connectBytes, _ := connectReq.MarshalBinary()

	buff := make([]byte, connectBufferSize)
	length := tr.request(connectBytes, buff)

	if length >= connectMinResponseLength {
		var response connectResponse
		response.UnmarshalBinary(buff[:length])

		if response.action == 3 || tr.transactionId != response.transactionId {
			error := errorResponse{}
			error.UnmarshalBinary(buff)
			return errors.New("Response Error: " + error.message)
		}

		tr.connectionId = response.connectionId
		return nil
	}

	return errors.New("Invalid connect response length")
}

func (tr *UdpTracker) Run() {
	log.Printf("Tracker : Run : Started (%s)\n", tr.announceURL)
	defer log.Printf("Tracker : Run : Completed (%s)\n", tr.announceURL)

	rand.Seed(time.Now().UnixNano())

	serverAddr, err := net.ResolveUDPAddr("udp", tr.announceURL.Host)
	if err != nil {
		log.Println("Tracker : Could not resolve tracker host!")
		return
	}

	conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: 0})
	if err != nil {
		log.Println("Tracker : Could not bind to any port")
		return
	}

	tr.transactionId = rand.Uint32()
	tr.serverAddr = serverAddr
	tr.conn = conn
	tr.Announce(Started)

	for {
		select {
		case <-tr.quit:
			log.Println("Tracker : Stop : Stopping")
			tr.Announce(Stopped)
			return
		case <-tr.completedCh:
			go tr.Announce(Completed)
		case <-tr.timer:
			log.Printf("Tracker : Run : Interval Timer Expired (%s)\n", tr.announceURL)
			go tr.Announce(Interval)
		case stats := <-tr.peerChans.stats:
			log.Println("read from stats", stats)
		}
	}
}

// Send a udp packet to the tracker, fill in the dest buffer with the response
func (tr *UdpTracker) request(payload []byte, dest []byte) int {
	n := 0
	totalAttempts := 0

	// notify the sender when a response is recieved
	recvChan := make(chan bool)

	// keep sending the packet and wait for the specifed time
	// before trying again, limit n to 8
	go func() {

		// initial send
		tr.conn.WriteTo(payload, tr.serverAddr)
	Listen:
		for {
			// timeout: 15 * 2 ^ n (0-8)
			timeout := time.Second * time.Duration(15 * int(math.Pow(2.0, float64(n))))
			timer := time.After(timeout)
			select {
			case <-recvChan:
				break Listen
			case <-timer:
				tr.conn.WriteTo(payload, tr.serverAddr)
				totalAttempts++
				if n < 8 { n++ }
			}
		}
	}()

	// read until some data is received
	var err error
	length := 0
	for length == 0 {
		length, _, err = tr.conn.ReadFrom(dest)

		// not sure what to do here?
		if err != nil {
			log.Println("Tracker : udp packet read error")
		}
	}

	recvChan <- true
	return length
}
<|endoftext|>