<commit_before>// Copyright 2009 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 rpc

import (
	"fmt"
	"http"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	serverAddr, newServerAddr string
	httpServerAddr            string
	once, newOnce, httpOnce   sync.Once
)

const (
	second      = 1e9
	newHttpPath = "/foo"
)

type Args struct {
	A, B int
}

type Reply struct {
	C int
}

type Arith int

func (t *Arith) Add(args *Args, reply *Reply) os.Error {
	reply.C = args.A + args.B
	return nil
}

func (t *Arith) Mul(args *Args, reply *Reply) os.Error {
	reply.C = args.A * args.B
	return nil
}

func (t *Arith) Div(args *Args, reply *Reply) os.Error {
	if args.B == 0 {
		return os.ErrorString("divide by zero")
	}
	reply.C = args.A / args.B
	return nil
}

func (t *Arith) String(args *Args, reply *string) os.Error {
	*reply = fmt.Sprintf("%d+%d=%d", args.A, args.B, args.A+args.B)
	return nil
}

func (t *Arith) Scan(args *string, reply *Reply) (err os.Error) {
	_, err = fmt.Sscan(*args, &reply.C)
	return
}

func (t *Arith) Error(args *Args, reply *Reply) os.Error {
	panic("ERROR")
}

func listenTCP() (net.Listener, string) {
	l, e := net.Listen("tcp", "127.0.0.1:0") // any available address
	if e != nil {
		log.Fatalf("net.Listen tcp :0: %v", e)
	}
	return l, l.Addr().String()
}

func startServer() {
	Register(new(Arith))

	var l net.Listener
	l, serverAddr = listenTCP()
	log.Println("Test RPC server listening on", serverAddr)
	go Accept(l)

	HandleHTTP()
	httpOnce.Do(startHttpServer)
}

func startNewServer() {
	s := NewServer()
	s.Register(new(Arith))

	var l net.Listener
	l, newServerAddr = listenTCP()
	log.Println("NewServer test RPC server listening on", newServerAddr)
	go Accept(l)

	s.HandleHTTP(newHttpPath, "/bar")
	httpOnce.Do(startHttpServer)
}

func startHttpServer() {
	var l net.Listener
	l, httpServerAddr = listenTCP()
	httpServerAddr = l.Addr().String()
	log.Println("Test HTTP RPC server listening on", httpServerAddr)
	go http.Serve(l, nil)
}

func TestRPC(t *testing.T) {
	once.Do(startServer)
	testRPC(t, serverAddr)
	newOnce.Do(startNewServer)
	testRPC(t, newServerAddr)
}

func testRPC(t *testing.T, addr string) {
	client, err := Dial("tcp", addr)
	if err != nil {
		t.Fatal("dialing", err)
	}

	// Synchronous calls
	args := &Args{7, 8}
	reply := new(Reply)
	err = client.Call("Arith.Add", args, reply)
	if err != nil {
		t.Errorf("Add: expected no error but got string %q", err.String())
	}
	if reply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", reply.C, args.A+args.B)
	}

	// Nonexistent method
	args = &Args{7, 0}
	reply = new(Reply)
	err = client.Call("Arith.BadOperation", args, reply)
	// expect an error
	if err == nil {
		t.Error("BadOperation: expected error")
	} else if !strings.HasPrefix(err.String(), "rpc: can't find method ") {
		t.Errorf("BadOperation: expected can't find method error; got %q", err)
	}

	// Unknown service
	args = &Args{7, 8}
	reply = new(Reply)
	err = client.Call("Arith.Unknown", args, reply)
	if err == nil {
		t.Error("expected error calling unknown service")
	} else if strings.Index(err.String(), "method") < 0 {
		t.Error("expected error about method; got", err)
	}

	// Out of order.
	args = &Args{7, 8}
	mulReply := new(Reply)
	mulCall := client.Go("Arith.Mul", args, mulReply, nil)
	addReply := new(Reply)
	addCall := client.Go("Arith.Add", args, addReply, nil)

	addCall = <-addCall.Done
	if addCall.Error != nil {
		t.Errorf("Add: expected no error but got string %q", addCall.Error.String())
	}
	if addReply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", addReply.C, args.A+args.B)
	}

	mulCall = <-mulCall.Done
	if mulCall.Error != nil {
		t.Errorf("Mul: expected no error but got string %q", mulCall.Error.String())
	}
	if mulReply.C != args.A*args.B {
		t.Errorf("Mul: expected %d got %d", mulReply.C, args.A*args.B)
	}

	// Error test
	args = &Args{7, 0}
	reply = new(Reply)
	err = client.Call("Arith.Div", args, reply)
	// expect an error: zero divide
	if err == nil {
		t.Error("Div: expected error")
	} else if err.String() != "divide by zero" {
		t.Error("Div: expected divide by zero error; got", err)
	}

	// Bad type.
	reply = new(Reply)
	err = client.Call("Arith.Add", reply, reply) // args, reply would be the correct thing to use
	if err == nil {
		t.Error("expected error calling Arith.Add with wrong arg type")
	} else if strings.Index(err.String(), "type") < 0 {
		t.Error("expected error about type; got", err)
	}

	// Non-struct argument
	const Val = 12345
	str := fmt.Sprint(Val)
	reply = new(Reply)
	err = client.Call("Arith.Scan", &str, reply)
	if err != nil {
		t.Errorf("Scan: expected no error but got string %q", err.String())
	} else if reply.C != Val {
		t.Errorf("Scan: expected %d got %d", Val, reply.C)
	}

	// Non-struct reply
	args = &Args{27, 35}
	str = ""
	err = client.Call("Arith.String", args, &str)
	if err != nil {
		t.Errorf("String: expected no error but got string %q", err.String())
	}
	expect := fmt.Sprintf("%d+%d=%d", args.A, args.B, args.A+args.B)
	if str != expect {
		t.Errorf("String: expected %s got %s", expect, str)
	}

	args = &Args{7, 8}
	reply = new(Reply)
	err = client.Call("Arith.Mul", args, reply)
	if err != nil {
		t.Errorf("Mul: expected no error but got string %q", err.String())
	}
	if reply.C != args.A*args.B {
		t.Errorf("Mul: expected %d got %d", reply.C, args.A*args.B)
	}
}

func TestHTTP(t *testing.T) {
	once.Do(startServer)
	testHTTPRPC(t, "")
	newOnce.Do(startNewServer)
	testHTTPRPC(t, newHttpPath)
}

func testHTTPRPC(t *testing.T, path string) {
	var client *Client
	var err os.Error
	if path == "" {
		client, err = DialHTTP("tcp", httpServerAddr)
	} else {
		client, err = DialHTTPPath("tcp", httpServerAddr, path)
	}
	if err != nil {
		t.Fatal("dialing", err)
	}

	// Synchronous calls
	args := &Args{7, 8}
	reply := new(Reply)
	err = client.Call("Arith.Add", args, reply)
	if err != nil {
		t.Errorf("Add: expected no error but got string %q", err.String())
	}
	if reply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", reply.C, args.A+args.B)
	}
}

type ArgNotPointer int
type ReplyNotPointer int
type ArgNotPublic int
type ReplyNotPublic int
type local struct{}

func (t *ArgNotPointer) ArgNotPointer(args Args, reply *Reply) os.Error {
	return nil
}

func (t *ReplyNotPointer) ReplyNotPointer(args *Args, reply Reply) os.Error {
	return nil
}

func (t *ArgNotPublic) ArgNotPublic(args *local, reply *Reply) os.Error {
	return nil
}

func (t *ReplyNotPublic) ReplyNotPublic(args *Args, reply *local) os.Error {
	return nil
}

// Check that registration handles lots of bad methods and a type with no suitable methods.
func TestRegistrationError(t *testing.T) {
	err := Register(new(ArgNotPointer))
	if err == nil {
		t.Errorf("expected error registering ArgNotPointer")
	}
	err = Register(new(ReplyNotPointer))
	if err == nil {
		t.Errorf("expected error registering ReplyNotPointer")
	}
	err = Register(new(ArgNotPublic))
	if err == nil {
		t.Errorf("expected error registering ArgNotPublic")
	}
	err = Register(new(ReplyNotPublic))
	if err == nil {
		t.Errorf("expected error registering ReplyNotPublic")
	}
}

type WriteFailCodec int

func (WriteFailCodec) WriteRequest(*Request, interface{}) os.Error {
	// the panic caused by this error used to not unlock a lock.
	return os.NewError("fail")
}

func (WriteFailCodec) ReadResponseHeader(*Response) os.Error {
	time.Sleep(60e9)
	panic("unreachable")
}

func (WriteFailCodec) ReadResponseBody(interface{}) os.Error {
	time.Sleep(60e9)
	panic("unreachable")
}

func (WriteFailCodec) Close() os.Error {
	return nil
}

func TestSendDeadlock(t *testing.T) {
	client := NewClientWithCodec(WriteFailCodec(0))

	done := make(chan bool)
	go func() {
		testSendDeadlock(client)
		testSendDeadlock(client)
		done <- true
	}()
	select {
	case <-done:
		return
	case <-time.After(5e9):
		t.Fatal("deadlock")
	}
}

func testSendDeadlock(client *Client) {
	defer func() {
		recover()
	}()
	args := &Args{7, 8}
	reply := new(Reply)
	client.Call("Arith.Add", args, reply)
}
<commit_msg>rpc: use httptest.Server for tests<commit_after>// Copyright 2009 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 rpc

import (
	"fmt"
	"http/httptest"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"testing"
	"time"
)

var (
	serverAddr, newServerAddr string
	httpServerAddr            string
	once, newOnce, httpOnce   sync.Once
)

const (
	second      = 1e9
	newHttpPath = "/foo"
)

type Args struct {
	A, B int
}

type Reply struct {
	C int
}

type Arith int

func (t *Arith) Add(args *Args, reply *Reply) os.Error {
	reply.C = args.A + args.B
	return nil
}

func (t *Arith) Mul(args *Args, reply *Reply) os.Error {
	reply.C = args.A * args.B
	return nil
}

func (t *Arith) Div(args *Args, reply *Reply) os.Error {
	if args.B == 0 {
		return os.ErrorString("divide by zero")
	}
	reply.C = args.A / args.B
	return nil
}

func (t *Arith) String(args *Args, reply *string) os.Error {
	*reply = fmt.Sprintf("%d+%d=%d", args.A, args.B, args.A+args.B)
	return nil
}

func (t *Arith) Scan(args *string, reply *Reply) (err os.Error) {
	_, err = fmt.Sscan(*args, &reply.C)
	return
}

func (t *Arith) Error(args *Args, reply *Reply) os.Error {
	panic("ERROR")
}

func listenTCP() (net.Listener, string) {
	l, e := net.Listen("tcp", "127.0.0.1:0") // any available address
	if e != nil {
		log.Fatalf("net.Listen tcp :0: %v", e)
	}
	return l, l.Addr().String()
}

func startServer() {
	Register(new(Arith))

	var l net.Listener
	l, serverAddr = listenTCP()
	log.Println("Test RPC server listening on", serverAddr)
	go Accept(l)

	HandleHTTP()
	httpOnce.Do(startHttpServer)
}

func startNewServer() {
	s := NewServer()
	s.Register(new(Arith))

	var l net.Listener
	l, newServerAddr = listenTCP()
	log.Println("NewServer test RPC server listening on", newServerAddr)
	go Accept(l)

	s.HandleHTTP(newHttpPath, "/bar")
	httpOnce.Do(startHttpServer)
}

func startHttpServer() {
	server := httptest.NewServer(nil)
	httpServerAddr = server.Listener.Addr().String()
	log.Println("Test HTTP RPC server listening on", httpServerAddr)
}

func TestRPC(t *testing.T) {
	once.Do(startServer)
	testRPC(t, serverAddr)
	newOnce.Do(startNewServer)
	testRPC(t, newServerAddr)
}

func testRPC(t *testing.T, addr string) {
	client, err := Dial("tcp", addr)
	if err != nil {
		t.Fatal("dialing", err)
	}

	// Synchronous calls
	args := &Args{7, 8}
	reply := new(Reply)
	err = client.Call("Arith.Add", args, reply)
	if err != nil {
		t.Errorf("Add: expected no error but got string %q", err.String())
	}
	if reply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", reply.C, args.A+args.B)
	}

	// Nonexistent method
	args = &Args{7, 0}
	reply = new(Reply)
	err = client.Call("Arith.BadOperation", args, reply)
	// expect an error
	if err == nil {
		t.Error("BadOperation: expected error")
	} else if !strings.HasPrefix(err.String(), "rpc: can't find method ") {
		t.Errorf("BadOperation: expected can't find method error; got %q", err)
	}

	// Unknown service
	args = &Args{7, 8}
	reply = new(Reply)
	err = client.Call("Arith.Unknown", args, reply)
	if err == nil {
		t.Error("expected error calling unknown service")
	} else if strings.Index(err.String(), "method") < 0 {
		t.Error("expected error about method; got", err)
	}

	// Out of order.
	args = &Args{7, 8}
	mulReply := new(Reply)
	mulCall := client.Go("Arith.Mul", args, mulReply, nil)
	addReply := new(Reply)
	addCall := client.Go("Arith.Add", args, addReply, nil)

	addCall = <-addCall.Done
	if addCall.Error != nil {
		t.Errorf("Add: expected no error but got string %q", addCall.Error.String())
	}
	if addReply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", addReply.C, args.A+args.B)
	}

	mulCall = <-mulCall.Done
	if mulCall.Error != nil {
		t.Errorf("Mul: expected no error but got string %q", mulCall.Error.String())
	}
	if mulReply.C != args.A*args.B {
		t.Errorf("Mul: expected %d got %d", mulReply.C, args.A*args.B)
	}

	// Error test
	args = &Args{7, 0}
	reply = new(Reply)
	err = client.Call("Arith.Div", args, reply)
	// expect an error: zero divide
	if err == nil {
		t.Error("Div: expected error")
	} else if err.String() != "divide by zero" {
		t.Error("Div: expected divide by zero error; got", err)
	}

	// Bad type.
	reply = new(Reply)
	err = client.Call("Arith.Add", reply, reply) // args, reply would be the correct thing to use
	if err == nil {
		t.Error("expected error calling Arith.Add with wrong arg type")
	} else if strings.Index(err.String(), "type") < 0 {
		t.Error("expected error about type; got", err)
	}

	// Non-struct argument
	const Val = 12345
	str := fmt.Sprint(Val)
	reply = new(Reply)
	err = client.Call("Arith.Scan", &str, reply)
	if err != nil {
		t.Errorf("Scan: expected no error but got string %q", err.String())
	} else if reply.C != Val {
		t.Errorf("Scan: expected %d got %d", Val, reply.C)
	}

	// Non-struct reply
	args = &Args{27, 35}
	str = ""
	err = client.Call("Arith.String", args, &str)
	if err != nil {
		t.Errorf("String: expected no error but got string %q", err.String())
	}
	expect := fmt.Sprintf("%d+%d=%d", args.A, args.B, args.A+args.B)
	if str != expect {
		t.Errorf("String: expected %s got %s", expect, str)
	}

	args = &Args{7, 8}
	reply = new(Reply)
	err = client.Call("Arith.Mul", args, reply)
	if err != nil {
		t.Errorf("Mul: expected no error but got string %q", err.String())
	}
	if reply.C != args.A*args.B {
		t.Errorf("Mul: expected %d got %d", reply.C, args.A*args.B)
	}
}

func TestHTTP(t *testing.T) {
	once.Do(startServer)
	testHTTPRPC(t, "")
	newOnce.Do(startNewServer)
	testHTTPRPC(t, newHttpPath)
}

func testHTTPRPC(t *testing.T, path string) {
	var client *Client
	var err os.Error
	if path == "" {
		client, err = DialHTTP("tcp", httpServerAddr)
	} else {
		client, err = DialHTTPPath("tcp", httpServerAddr, path)
	}
	if err != nil {
		t.Fatal("dialing", err)
	}

	// Synchronous calls
	args := &Args{7, 8}
	reply := new(Reply)
	err = client.Call("Arith.Add", args, reply)
	if err != nil {
		t.Errorf("Add: expected no error but got string %q", err.String())
	}
	if reply.C != args.A+args.B {
		t.Errorf("Add: expected %d got %d", reply.C, args.A+args.B)
	}
}

type ArgNotPointer int
type ReplyNotPointer int
type ArgNotPublic int
type ReplyNotPublic int
type local struct{}

func (t *ArgNotPointer) ArgNotPointer(args Args, reply *Reply) os.Error {
	return nil
}

func (t *ReplyNotPointer) ReplyNotPointer(args *Args, reply Reply) os.Error {
	return nil
}

func (t *ArgNotPublic) ArgNotPublic(args *local, reply *Reply) os.Error {
	return nil
}

func (t *ReplyNotPublic) ReplyNotPublic(args *Args, reply *local) os.Error {
	return nil
}

// Check that registration handles lots of bad methods and a type with no suitable methods.
func TestRegistrationError(t *testing.T) {
	err := Register(new(ArgNotPointer))
	if err == nil {
		t.Errorf("expected error registering ArgNotPointer")
	}
	err = Register(new(ReplyNotPointer))
	if err == nil {
		t.Errorf("expected error registering ReplyNotPointer")
	}
	err = Register(new(ArgNotPublic))
	if err == nil {
		t.Errorf("expected error registering ArgNotPublic")
	}
	err = Register(new(ReplyNotPublic))
	if err == nil {
		t.Errorf("expected error registering ReplyNotPublic")
	}
}

type WriteFailCodec int

func (WriteFailCodec) WriteRequest(*Request, interface{}) os.Error {
	// the panic caused by this error used to not unlock a lock.
	return os.NewError("fail")
}

func (WriteFailCodec) ReadResponseHeader(*Response) os.Error {
	time.Sleep(60e9)
	panic("unreachable")
}

func (WriteFailCodec) ReadResponseBody(interface{}) os.Error {
	time.Sleep(60e9)
	panic("unreachable")
}

func (WriteFailCodec) Close() os.Error {
	return nil
}

func TestSendDeadlock(t *testing.T) {
	client := NewClientWithCodec(WriteFailCodec(0))

	done := make(chan bool)
	go func() {
		testSendDeadlock(client)
		testSendDeadlock(client)
		done <- true
	}()
	select {
	case <-done:
		return
	case <-time.After(5e9):
		t.Fatal("deadlock")
	}
}

func testSendDeadlock(client *Client) {
	defer func() {
		recover()
	}()
	args := &Args{7, 8}
	reply := new(Reply)
	client.Call("Arith.Add", args, reply)
}
<|endoftext|>


<commit_before>package torrent

import (
	"container/heap"
	"context"
	"encoding/gob"
	"reflect"
	"runtime/pprof"
	"time"
	"unsafe"

	"github.com/anacrolix/log"
	"github.com/anacrolix/multiless"

	request_strategy "github.com/anacrolix/torrent/request-strategy"
)

func (t *Torrent) requestStrategyPieceOrderState(i int) request_strategy.PieceRequestOrderState {
	return request_strategy.PieceRequestOrderState{
		Priority:     t.piece(i).purePriority(),
		Partial:      t.piecePartiallyDownloaded(i),
		Availability: t.piece(i).availability(),
	}
}

func init() {
	gob.Register(peerId{})
}

type peerId struct {
	*Peer
	ptr uintptr
}

func (p peerId) Uintptr() uintptr {
	return p.ptr
}

func (p peerId) GobEncode() (b []byte, _ error) {
	*(*reflect.SliceHeader)(unsafe.Pointer(&b)) = reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(&p.ptr)),
		Len:  int(unsafe.Sizeof(p.ptr)),
		Cap:  int(unsafe.Sizeof(p.ptr)),
	}
	return
}

func (p *peerId) GobDecode(b []byte) error {
	if uintptr(len(b)) != unsafe.Sizeof(p.ptr) {
		panic(len(b))
	}
	ptr := unsafe.Pointer(&b[0])
	p.ptr = *(*uintptr)(ptr)
	log.Printf("%p", ptr)
	dst := reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(&p.Peer)),
		Len:  int(unsafe.Sizeof(p.Peer)),
		Cap:  int(unsafe.Sizeof(p.Peer)),
	}
	copy(*(*[]byte)(unsafe.Pointer(&dst)), b)
	return nil
}

type (
	RequestIndex   = request_strategy.RequestIndex
	chunkIndexType = request_strategy.ChunkIndex
)

type peerRequests struct {
	requestIndexes []RequestIndex
	peer           *Peer
}

func (p *peerRequests) Len() int {
	return len(p.requestIndexes)
}

func (p *peerRequests) Less(i, j int) bool {
	leftRequest := p.requestIndexes[i]
	rightRequest := p.requestIndexes[j]
	t := p.peer.t
	leftPieceIndex := leftRequest / t.chunksPerRegularPiece()
	rightPieceIndex := rightRequest / t.chunksPerRegularPiece()
	ml := multiless.New()
	// Push requests that can't be served right now to the end. But we don't throw them away unless
	// there's a better alternative. This is for when we're using the fast extension and get choked
	// but our requests could still be good when we get unchoked.
	if p.peer.peerChoking {
		ml = ml.Bool(
			!p.peer.peerAllowedFast.Contains(leftPieceIndex),
			!p.peer.peerAllowedFast.Contains(rightPieceIndex),
		)
	}
	leftPeer := t.pendingRequests[leftRequest]
	rightPeer := t.pendingRequests[rightRequest]
	ml = ml.Bool(rightPeer == p.peer, leftPeer == p.peer)
	ml = ml.Bool(rightPeer == nil, leftPeer == nil)
	if ml.Ok() {
		return ml.MustLess()
	}
	if leftPeer != nil {
		// The right peer should also be set, or we'd have resolved the computation by now.
		ml = ml.Uint64(
			rightPeer.requestState.Requests.GetCardinality(),
			leftPeer.requestState.Requests.GetCardinality(),
		)
		// Could either of the lastRequested be Zero? That's what checking an existing peer is for.
		leftLast := t.lastRequested[leftRequest]
		rightLast := t.lastRequested[rightRequest]
		if leftLast.IsZero() || rightLast.IsZero() {
			panic("expected non-zero last requested times")
		}
		// We want the most-recently requested on the left. Clients like Transmission serve requests
		// in received order, so the most recently-requested is the one that has the longest until
		// it will be served and therefore is the best candidate to cancel.
		ml = ml.CmpInt64(rightLast.Sub(leftLast).Nanoseconds())
	}
	leftPiece := t.piece(int(leftPieceIndex))
	rightPiece := t.piece(int(rightPieceIndex))
	ml = ml.Int(
		// Technically we would be happy with the cached priority here, except we don't actually
		// cache it anymore, and Torrent.piecePriority just does another lookup of *Piece to resolve
		// the priority through Piece.purePriority, which is probably slower.
		-int(leftPiece.purePriority()),
		-int(rightPiece.purePriority()),
	)
	ml = ml.Int(
		int(leftPiece.relativeAvailability),
		int(rightPiece.relativeAvailability))
	return ml.Less()
}

func (p *peerRequests) Swap(i, j int) {
	p.requestIndexes[i], p.requestIndexes[j] = p.requestIndexes[j], p.requestIndexes[i]
}

func (p *peerRequests) Push(x interface{}) {
	p.requestIndexes = append(p.requestIndexes, x.(RequestIndex))
}

func (p *peerRequests) Pop() interface{} {
	last := len(p.requestIndexes) - 1
	x := p.requestIndexes[last]
	p.requestIndexes = p.requestIndexes[:last]
	return x
}

type desiredRequestState struct {
	Requests   peerRequests
	Interested bool
}

func (p *Peer) getDesiredRequestState() (desired desiredRequestState) {
	if !p.t.haveInfo() {
		return
	}
	if p.t.closed.IsSet() {
		return
	}
	input := p.t.getRequestStrategyInput()
	requestHeap := peerRequests{
		peer: p,
	}
	request_strategy.GetRequestablePieces(
		input,
		p.t.getPieceRequestOrder(),
		func(ih InfoHash, pieceIndex int) {
			if ih != p.t.infoHash {
				return
			}
			if !p.peerHasPiece(pieceIndex) {
				return
			}
			allowedFast := p.peerAllowedFast.ContainsInt(pieceIndex)
			p.t.piece(pieceIndex).undirtiedChunksIter.Iter(func(ci request_strategy.ChunkIndex) {
				r := p.t.pieceRequestIndexOffset(pieceIndex) + ci
				if !allowedFast {
					// We must signal interest to request this. TODO: We could set interested if the
					// peers pieces (minus the allowed fast set) overlap with our missing pieces if
					// there are any readers, or any pending pieces.
					desired.Interested = true
					// We can make or will allow sustaining a request here if we're not choked, or
					// have made the request previously (presumably while unchoked), and haven't had
					// the peer respond yet (and the request was retained because we are using the
					// fast extension).
					if p.peerChoking && !p.requestState.Requests.Contains(r) {
						// We can't request this right now.
						return
					}
				}
				if p.requestState.Cancelled.Contains(r) {
					// Can't re-request while awaiting acknowledgement.
					return
				}
				requestHeap.requestIndexes = append(requestHeap.requestIndexes, r)
			})
		},
	)
	p.t.assertPendingRequests()
	desired.Requests = requestHeap
	return
}

func (p *Peer) maybeUpdateActualRequestState() bool {
	if p.needRequestUpdate == "" {
		return true
	}
	var more bool
	pprof.Do(
		context.Background(),
		pprof.Labels("update request", p.needRequestUpdate),
		func(_ context.Context) {
			next := p.getDesiredRequestState()
			more = p.applyRequestState(next)
		},
	)
	return more
}

// Transmit/action the request state to the peer.
func (p *Peer) applyRequestState(next desiredRequestState) bool {
	current := &p.requestState
	if !p.setInterested(next.Interested) {
		return false
	}
	more := true
	requestHeap := &next.Requests
	t := p.t
	heap.Init(requestHeap)
	for requestHeap.Len() != 0 && maxRequests(current.Requests.GetCardinality()) < p.nominalMaxRequests() {
		req := heap.Pop(requestHeap).(RequestIndex)
		existing := t.requestingPeer(req)
		if existing != nil && existing != p {
			// Don't steal from the poor.
			diff := int64(current.Requests.GetCardinality()) + 1 - (int64(existing.uncancelledRequests()) - 1)
			// Steal a request that leaves us with one more request than the existing peer
			// connection if the stealer more recently received a chunk.
			if diff > 1 || (diff == 1 && p.lastUsefulChunkReceived.Before(existing.lastUsefulChunkReceived)) {
				continue
			}
			t.cancelRequest(req)
		}
		more = p.mustRequest(req)
		if !more {
			break
		}
	}
	// TODO: This may need to change, we might want to update even if there were no requests due to
	// filtering them for being recently requested already.
	p.updateRequestsTimer.Stop()
	if more {
		p.needRequestUpdate = ""
		if current.Interested {
			p.updateRequestsTimer.Reset(3 * time.Second)
		}
	}
	return more
}
<commit_msg>Include requests pending cancel in current request count<commit_after>package torrent

import (
	"container/heap"
	"context"
	"encoding/gob"
	"reflect"
	"runtime/pprof"
	"time"
	"unsafe"

	"github.com/anacrolix/log"
	"github.com/anacrolix/multiless"

	request_strategy "github.com/anacrolix/torrent/request-strategy"
)

func (t *Torrent) requestStrategyPieceOrderState(i int) request_strategy.PieceRequestOrderState {
	return request_strategy.PieceRequestOrderState{
		Priority:     t.piece(i).purePriority(),
		Partial:      t.piecePartiallyDownloaded(i),
		Availability: t.piece(i).availability(),
	}
}

func init() {
	gob.Register(peerId{})
}

type peerId struct {
	*Peer
	ptr uintptr
}

func (p peerId) Uintptr() uintptr {
	return p.ptr
}

func (p peerId) GobEncode() (b []byte, _ error) {
	*(*reflect.SliceHeader)(unsafe.Pointer(&b)) = reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(&p.ptr)),
		Len:  int(unsafe.Sizeof(p.ptr)),
		Cap:  int(unsafe.Sizeof(p.ptr)),
	}
	return
}

func (p *peerId) GobDecode(b []byte) error {
	if uintptr(len(b)) != unsafe.Sizeof(p.ptr) {
		panic(len(b))
	}
	ptr := unsafe.Pointer(&b[0])
	p.ptr = *(*uintptr)(ptr)
	log.Printf("%p", ptr)
	dst := reflect.SliceHeader{
		Data: uintptr(unsafe.Pointer(&p.Peer)),
		Len:  int(unsafe.Sizeof(p.Peer)),
		Cap:  int(unsafe.Sizeof(p.Peer)),
	}
	copy(*(*[]byte)(unsafe.Pointer(&dst)), b)
	return nil
}

type (
	RequestIndex   = request_strategy.RequestIndex
	chunkIndexType = request_strategy.ChunkIndex
)

type peerRequests struct {
	requestIndexes []RequestIndex
	peer           *Peer
}

func (p *peerRequests) Len() int {
	return len(p.requestIndexes)
}

func (p *peerRequests) Less(i, j int) bool {
	leftRequest := p.requestIndexes[i]
	rightRequest := p.requestIndexes[j]
	t := p.peer.t
	leftPieceIndex := leftRequest / t.chunksPerRegularPiece()
	rightPieceIndex := rightRequest / t.chunksPerRegularPiece()
	ml := multiless.New()
	// Push requests that can't be served right now to the end. But we don't throw them away unless
	// there's a better alternative. This is for when we're using the fast extension and get choked
	// but our requests could still be good when we get unchoked.
	if p.peer.peerChoking {
		ml = ml.Bool(
			!p.peer.peerAllowedFast.Contains(leftPieceIndex),
			!p.peer.peerAllowedFast.Contains(rightPieceIndex),
		)
	}
	leftPeer := t.pendingRequests[leftRequest]
	rightPeer := t.pendingRequests[rightRequest]
	ml = ml.Bool(rightPeer == p.peer, leftPeer == p.peer)
	ml = ml.Bool(rightPeer == nil, leftPeer == nil)
	if ml.Ok() {
		return ml.MustLess()
	}
	if leftPeer != nil {
		// The right peer should also be set, or we'd have resolved the computation by now.
		ml = ml.Uint64(
			rightPeer.requestState.Requests.GetCardinality(),
			leftPeer.requestState.Requests.GetCardinality(),
		)
		// Could either of the lastRequested be Zero? That's what checking an existing peer is for.
		leftLast := t.lastRequested[leftRequest]
		rightLast := t.lastRequested[rightRequest]
		if leftLast.IsZero() || rightLast.IsZero() {
			panic("expected non-zero last requested times")
		}
		// We want the most-recently requested on the left. Clients like Transmission serve requests
		// in received order, so the most recently-requested is the one that has the longest until
		// it will be served and therefore is the best candidate to cancel.
		ml = ml.CmpInt64(rightLast.Sub(leftLast).Nanoseconds())
	}
	leftPiece := t.piece(int(leftPieceIndex))
	rightPiece := t.piece(int(rightPieceIndex))
	ml = ml.Int(
		// Technically we would be happy with the cached priority here, except we don't actually
		// cache it anymore, and Torrent.piecePriority just does another lookup of *Piece to resolve
		// the priority through Piece.purePriority, which is probably slower.
		-int(leftPiece.purePriority()),
		-int(rightPiece.purePriority()),
	)
	ml = ml.Int(
		int(leftPiece.relativeAvailability),
		int(rightPiece.relativeAvailability))
	return ml.Less()
}

func (p *peerRequests) Swap(i, j int) {
	p.requestIndexes[i], p.requestIndexes[j] = p.requestIndexes[j], p.requestIndexes[i]
}

func (p *peerRequests) Push(x interface{}) {
	p.requestIndexes = append(p.requestIndexes, x.(RequestIndex))
}

func (p *peerRequests) Pop() interface{} {
	last := len(p.requestIndexes) - 1
	x := p.requestIndexes[last]
	p.requestIndexes = p.requestIndexes[:last]
	return x
}

type desiredRequestState struct {
	Requests   peerRequests
	Interested bool
}

func (p *Peer) getDesiredRequestState() (desired desiredRequestState) {
	if !p.t.haveInfo() {
		return
	}
	if p.t.closed.IsSet() {
		return
	}
	input := p.t.getRequestStrategyInput()
	requestHeap := peerRequests{
		peer: p,
	}
	request_strategy.GetRequestablePieces(
		input,
		p.t.getPieceRequestOrder(),
		func(ih InfoHash, pieceIndex int) {
			if ih != p.t.infoHash {
				return
			}
			if !p.peerHasPiece(pieceIndex) {
				return
			}
			allowedFast := p.peerAllowedFast.ContainsInt(pieceIndex)
			p.t.piece(pieceIndex).undirtiedChunksIter.Iter(func(ci request_strategy.ChunkIndex) {
				r := p.t.pieceRequestIndexOffset(pieceIndex) + ci
				if !allowedFast {
					// We must signal interest to request this. TODO: We could set interested if the
					// peers pieces (minus the allowed fast set) overlap with our missing pieces if
					// there are any readers, or any pending pieces.
					desired.Interested = true
					// We can make or will allow sustaining a request here if we're not choked, or
					// have made the request previously (presumably while unchoked), and haven't had
					// the peer respond yet (and the request was retained because we are using the
					// fast extension).
					if p.peerChoking && !p.requestState.Requests.Contains(r) {
						// We can't request this right now.
						return
					}
				}
				if p.requestState.Cancelled.Contains(r) {
					// Can't re-request while awaiting acknowledgement.
					return
				}
				requestHeap.requestIndexes = append(requestHeap.requestIndexes, r)
			})
		},
	)
	p.t.assertPendingRequests()
	desired.Requests = requestHeap
	return
}

func (p *Peer) maybeUpdateActualRequestState() bool {
	if p.needRequestUpdate == "" {
		return true
	}
	var more bool
	pprof.Do(
		context.Background(),
		pprof.Labels("update request", p.needRequestUpdate),
		func(_ context.Context) {
			next := p.getDesiredRequestState()
			more = p.applyRequestState(next)
		},
	)
	return more
}

// Transmit/action the request state to the peer.
func (p *Peer) applyRequestState(next desiredRequestState) bool {
	current := &p.requestState
	if !p.setInterested(next.Interested) {
		return false
	}
	more := true
	requestHeap := &next.Requests
	t := p.t
	heap.Init(requestHeap)
	for requestHeap.Len() != 0 && maxRequests(current.Requests.GetCardinality()+current.Cancelled.GetCardinality()) < p.nominalMaxRequests() {
		req := heap.Pop(requestHeap).(RequestIndex)
		existing := t.requestingPeer(req)
		if existing != nil && existing != p {
			// Don't steal from the poor.
			diff := int64(current.Requests.GetCardinality()) + 1 - (int64(existing.uncancelledRequests()) - 1)
			// Steal a request that leaves us with one more request than the existing peer
			// connection if the stealer more recently received a chunk.
			if diff > 1 || (diff == 1 && p.lastUsefulChunkReceived.Before(existing.lastUsefulChunkReceived)) {
				continue
			}
			t.cancelRequest(req)
		}
		more = p.mustRequest(req)
		if !more {
			break
		}
	}
	// TODO: This may need to change, we might want to update even if there were no requests due to
	// filtering them for being recently requested already.
	p.updateRequestsTimer.Stop()
	if more {
		p.needRequestUpdate = ""
		if current.Interested {
			p.updateRequestsTimer.Reset(3 * time.Second)
		}
	}
	return more
}
<|endoftext|>


<commit_before>// Package replay handles the serving of recorded data through a HTTP router.
package replay

import (
	"log"
	"net/http"
	"strconv"

	"github.com/1lann/lol-replay/record"
	"github.com/1lann/lol-replay/recording"
	"github.com/julienschmidt/httprouter"
)

// PathHeader is the common path header used for requests to the spectator
// endpoint.
const PathHeader = "/observer-mode/rest/consumer"

// A Retriever provides a recording for a given game ID and region.
// A nil recording should be returned if the recording does not exist.
type Retriever func(region, gameId string) *recording.Recording

type requestHandler struct {
	retriever Retriever
}

type httpWriterPipe struct {
	w           http.ResponseWriter
	contentType string
	hasWritten  bool
}

func newHTTPWriterPipe(w http.ResponseWriter,
	contentType string) *httpWriterPipe {
	return &httpWriterPipe{w: w, contentType: contentType, hasWritten: false}
}

func (p *httpWriterPipe) Write(data []byte) (int, error) {
	if !p.hasWritten {
		p.w.Header().Set("Content-Type", p.contentType)
		p.w.WriteHeader(http.StatusOK)
		p.hasWritten = true
	}

	return p.w.Write(data)
}

func (p *httpWriterPipe) HasWritten() bool {
	return p.hasWritten
}

func (rh requestHandler) version(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	version, err := record.GetPlatformVersion("OC1")
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("version unavailable"))
		return
	}

	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(version))
	return
}

func (rh requestHandler) getGameMetadata(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/json")
	_, err := rec.RetrieveGameMetadataTo(pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getGameMetadata silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("metadata not found"))
			return
		}

		log.Println("getGameMetadata error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

func (rh requestHandler) getLastChunkInfo(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	if ps.ByName("end") == "0" {
		rec.RetrieveLastChunkInfo().WriteTo(w)
	} else {
		rec.RetrieveFirstChunkInfo().WriteTo(w)
	}
}

func (rh requestHandler) getGameDataChunk(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	chunk, err := strconv.Atoi(ps.ByName("chunk"))
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("invalid chunk number"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/octet-stream")
	_, err = rec.RetrieveChunkTo(chunk, pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getGameDataChunk silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("chunk not found"))
			return
		}

		log.Println("getGameDataChunk error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

func (rh requestHandler) getKeyFrame(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	frame, err := strconv.Atoi(ps.ByName("frame"))
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("invalid keyframe number"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/octet-stream")
	_, err = rec.RetrieveKeyFrameTo(frame, pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getKeyFrame silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("keyframe not found"))
			return
		}

		log.Println("getKeyFrame error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

// Router returns a http.Handler that handles requests for recorded data.
func Router(retriever Retriever) http.Handler {
	handler := requestHandler{retriever}

	router := httprouter.New()
	router.GET(PathHeader+"/version", handler.version)
	router.GET(PathHeader+"/getGameMetaData/:region/:id/*ignore",
		handler.getGameMetadata)
	router.GET(PathHeader+"/getLastChunkInfo/:region/:id/:end/*ignore",
		handler.getLastChunkInfo)
	router.GET(PathHeader+"/getGameDataChunk/:region/:id/:chunk/*ignore",
		handler.getGameDataChunk)
	router.GET(PathHeader+"/getKeyFrame/:region/:id/:frame/*ignore",
		handler.getKeyFrame)

	return router
}
<commit_msg>Update the getLastChunkInfo endpoint so that it lets the spectator spectate from the beginning.<commit_after>// Package replay handles the serving of recorded data through a HTTP router.
package replay

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/1lann/lol-replay/record"
	"github.com/1lann/lol-replay/recording"
	"github.com/Clever/leakybucket"
	memorybucket "github.com/Clever/leakybucket/memory"
	"github.com/julienschmidt/httprouter"
)

// PathHeader is the common path header used for requests to the spectator
// endpoint.
const PathHeader = "/observer-mode/rest/consumer"

// A client is identified by a IP/gameID pair.
// We want to identify a client because in order to start a spectating session from the
// beginning, we need to "pretend" that the last available chunk is one of the first chunks.
// Therefore if the client is new, we modify the behaviour of getLastChunkInfo.
type client struct {
	IP     string
	gameID string
}

var bucketStore = memorybucket.New()

// A Retriever provides a recording for a given game ID and region.
// A nil recording should be returned if the recording does not exist.
type Retriever func(region, gameId string) *recording.Recording

type requestHandler struct {
	retriever Retriever
	// Track how often a client has been making requests, using the leaky
	// bucket algorithm so that if the same client spectates again after a while,
	// we consider them a new client.
	newClientBuckets map[client]leakybucket.Bucket
}

type httpWriterPipe struct {
	w           http.ResponseWriter
	contentType string
	hasWritten  bool
}

func newHTTPWriterPipe(w http.ResponseWriter,
	contentType string) *httpWriterPipe {
	return &httpWriterPipe{w: w, contentType: contentType, hasWritten: false}
}

func (p *httpWriterPipe) Write(data []byte) (int, error) {
	if !p.hasWritten {
		p.w.Header().Set("Content-Type", p.contentType)
		p.w.WriteHeader(http.StatusOK)
		p.hasWritten = true
	}

	return p.w.Write(data)
}

func (p *httpWriterPipe) HasWritten() bool {
	return p.hasWritten
}

func (rh requestHandler) version(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	version, err := record.GetPlatformVersion("OC1")
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("version unavailable"))
		return
	}

	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(version))
	return
}

func (rh requestHandler) getGameMetadata(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/json")
	_, err := rec.RetrieveGameMetadataTo(pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getGameMetadata silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("metadata not found"))
			return
		}

		log.Println("getGameMetadata error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

func (rh requestHandler) getLastChunkInfo(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	// Identify the client by the IP/gameID tuple
	ip, _, _ := net.SplitHostPort(r.RemoteAddr)
	c := client{
		IP:     ip,
		gameID: ps.ByName("id"),
	}

	if rh.newClientBuckets[c] == nil {
		// A normal spectator client should make request to this endpoint once every 10 seconds.
		// Therefore, we use the more conservative number of 3 here, meaning that a client is
		// considered new if it hasn't made 3 requests in the last minute.
		bucket, err := bucketStore.Create(
			fmt.Sprintf("%s-%s", r.RemoteAddr, ps.ByName("id")),
			3,
			time.Minute,
		)
		if err != nil {
			rec.RetrieveLastChunkInfo().WriteTo(w)
			return
		}
		rh.newClientBuckets[c] = bucket
	}

	// We try to figure out if the client is a new or not.  If the client is new, we "pretend"
	// that the last available chunk is one of the first few chunks, so that the spectator client
	// would start playing from the beginning.  Otherwise, we return the real last available chunk.
	_, err := rh.newClientBuckets[c].Add(1)
	if err != nil {
		rec.RetrieveLastChunkInfo().WriteTo(w)
	} else {
		rec.RetrieveFirstChunkInfo().WriteTo(w)
	}
}

func (rh requestHandler) getGameDataChunk(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	chunk, err := strconv.Atoi(ps.ByName("chunk"))
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("invalid chunk number"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/octet-stream")
	_, err = rec.RetrieveChunkTo(chunk, pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getGameDataChunk silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("chunk not found"))
			return
		}

		log.Println("getGameDataChunk error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

func (rh requestHandler) getKeyFrame(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	rec := rh.retriever(ps.ByName("region"), ps.ByName("id"))
	if rec == nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("game not found"))
		return
	}

	frame, err := strconv.Atoi(ps.ByName("frame"))
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("invalid keyframe number"))
		return
	}

	pipe := newHTTPWriterPipe(w, "application/octet-stream")
	_, err = rec.RetrieveKeyFrameTo(frame, pipe)
	if err != nil {
		if pipe.HasWritten() {
			log.Println("getKeyFrame silent error:", err)
			return
		}

		if err == recording.ErrMissingData {
			w.Header().Set("Content-Type", "text/plain")
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("keyframe not found"))
			return
		}

		log.Println("getKeyFrame error:", err)

		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("internal server error"))
		return
	}
}

// Router returns a http.Handler that handles requests for recorded data.
func Router(retriever Retriever) http.Handler {
	handler := requestHandler{
		retriever:        retriever,
		newClientBuckets: make(map[client]leakybucket.Bucket),
	}

	router := httprouter.New()
	router.GET(PathHeader+"/version", handler.version)
	router.GET(PathHeader+"/getGameMetaData/:region/:id/*ignore",
		handler.getGameMetadata)
	router.GET(PathHeader+"/getLastChunkInfo/:region/:id/:end/*ignore",
		handler.getLastChunkInfo)
	router.GET(PathHeader+"/getGameDataChunk/:region/:id/:chunk/*ignore",
		handler.getGameDataChunk)
	router.GET(PathHeader+"/getKeyFrame/:region/:id/:frame/*ignore",
		handler.getKeyFrame)

	return router
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"text/template"

	"github.com/docker/docker/api/types"
	flag "github.com/docker/docker/pkg/mflag"
)

// CmdInspect displays low-level information on one or more containers or images.
//
// Usage: docker inspect [OPTIONS] CONTAINER|IMAGE [CONTAINER|IMAGE...]

func (cli *DockerCli) CmdInspect(args ...string) error {
	cmd := cli.Subcmd("inspect", "CONTAINER|IMAGE [CONTAINER|IMAGE...]", "Return low-level information on a container or image", true)
	tmplStr := cmd.String([]string{"f", "#format", "-format"}, "", "Format the output using the given go template")
	cmd.Require(flag.Min, 1)

	cmd.ParseFlags(args, true)

	var tmpl *template.Template
	if *tmplStr != "" {
		var err error
		if tmpl, err = template.New("").Funcs(funcMap).Parse(*tmplStr); err != nil {
			fmt.Fprintf(cli.err, "Template parsing error: %v\n", err)
			return StatusError{StatusCode: 64,
				Status: "Template parsing error: " + err.Error()}
		}
	}

	indented := new(bytes.Buffer)
	indented.WriteByte('[')
	status := 0
	isImage := false

	for _, name := range cmd.Args() {
		obj, _, err := readBody(cli.call("GET", "/containers/"+name+"/json", nil, nil))
		if err != nil {
			obj, _, err = readBody(cli.call("GET", "/images/"+name+"/json", nil, nil))
			isImage = true
			if err != nil {
				if strings.Contains(err.Error(), "No such") {
					fmt.Fprintf(cli.err, "Error: No such image or container: %s\n", name)
				} else {
					fmt.Fprintf(cli.err, "%s", err)
				}
				status = 1
				continue
			}
		}

		if tmpl == nil {
			if err = json.Indent(indented, obj, "", "    "); err != nil {
				fmt.Fprintf(cli.err, "%s\n", err)
				status = 1
				continue
			}
		} else {
			dec := json.NewDecoder(bytes.NewReader(obj))

			if isImage {
				inspPtr := types.ImageInspect{}
				if err := dec.Decode(&inspPtr); err != nil {
					fmt.Fprintf(cli.err, "%s\n", err)
					status = 1
					continue
				}
				if err := tmpl.Execute(cli.out, inspPtr); err != nil {
					return err
				}
			} else {
				inspPtr := types.ContainerJSON{}
				if err := dec.Decode(&inspPtr); err != nil {
					fmt.Fprintf(cli.err, "%s\n", err)
					status = 1
					continue
				}
				if err := tmpl.Execute(cli.out, inspPtr); err != nil {
					return err

				}
			}
			cli.out.Write([]byte{'\n'})
		}
		indented.WriteString(",")
	}

	if indented.Len() > 1 {
		// Remove trailing ','
		indented.Truncate(indented.Len() - 1)
	}
	indented.WriteString("]\n")

	if tmpl == nil {
		if _, err := io.Copy(cli.out, indented); err != nil {
			return err
		}
	}

	if status != 0 {
		return StatusError{StatusCode: status}
	}
	return nil
}
<commit_msg> Remove empty line after client.CmdInspect docstring #12706<commit_after>package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"text/template"

	"github.com/docker/docker/api/types"
	flag "github.com/docker/docker/pkg/mflag"
)

// CmdInspect displays low-level information on one or more containers or images.
//
// Usage: docker inspect [OPTIONS] CONTAINER|IMAGE [CONTAINER|IMAGE...]
func (cli *DockerCli) CmdInspect(args ...string) error {
	cmd := cli.Subcmd("inspect", "CONTAINER|IMAGE [CONTAINER|IMAGE...]", "Return low-level information on a container or image", true)
	tmplStr := cmd.String([]string{"f", "#format", "-format"}, "", "Format the output using the given go template")
	cmd.Require(flag.Min, 1)

	cmd.ParseFlags(args, true)

	var tmpl *template.Template
	if *tmplStr != "" {
		var err error
		if tmpl, err = template.New("").Funcs(funcMap).Parse(*tmplStr); err != nil {
			fmt.Fprintf(cli.err, "Template parsing error: %v\n", err)
			return StatusError{StatusCode: 64,
				Status: "Template parsing error: " + err.Error()}
		}
	}

	indented := new(bytes.Buffer)
	indented.WriteByte('[')
	status := 0
	isImage := false

	for _, name := range cmd.Args() {
		obj, _, err := readBody(cli.call("GET", "/containers/"+name+"/json", nil, nil))
		if err != nil {
			obj, _, err = readBody(cli.call("GET", "/images/"+name+"/json", nil, nil))
			isImage = true
			if err != nil {
				if strings.Contains(err.Error(), "No such") {
					fmt.Fprintf(cli.err, "Error: No such image or container: %s\n", name)
				} else {
					fmt.Fprintf(cli.err, "%s", err)
				}
				status = 1
				continue
			}
		}

		if tmpl == nil {
			if err = json.Indent(indented, obj, "", "    "); err != nil {
				fmt.Fprintf(cli.err, "%s\n", err)
				status = 1
				continue
			}
		} else {
			dec := json.NewDecoder(bytes.NewReader(obj))

			if isImage {
				inspPtr := types.ImageInspect{}
				if err := dec.Decode(&inspPtr); err != nil {
					fmt.Fprintf(cli.err, "%s\n", err)
					status = 1
					continue
				}
				if err := tmpl.Execute(cli.out, inspPtr); err != nil {
					return err
				}
			} else {
				inspPtr := types.ContainerJSON{}
				if err := dec.Decode(&inspPtr); err != nil {
					fmt.Fprintf(cli.err, "%s\n", err)
					status = 1
					continue
				}
				if err := tmpl.Execute(cli.out, inspPtr); err != nil {
					return err

				}
			}
			cli.out.Write([]byte{'\n'})
		}
		indented.WriteString(",")
	}

	if indented.Len() > 1 {
		// Remove trailing ','
		indented.Truncate(indented.Len() - 1)
	}
	indented.WriteString("]\n")

	if tmpl == nil {
		if _, err := io.Copy(cli.out, indented); err != nil {
			return err
		}
	}

	if status != 0 {
		return StatusError{StatusCode: status}
	}
	return nil
}
<|endoftext|>


<commit_before><commit_msg>make code more robust<commit_after><|endoftext|>


<commit_before>package pingbeat

import (
	"errors"
	// "github.com/davecgh/go-spew/spew"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/cfgfile"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	cfg "github.com/joshuar/pingbeat/config"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"gopkg.in/go-playground/pool.v3"
	"net"
	"os"
	"sync"
	"time"
)

// Pingbeat struct contains all options and
// hosts to ping
type Pingbeat struct {
	isAlive     bool
	useIPv4     bool
	useIPv6     bool
	period      time.Duration
	timeout     time.Duration
	ipv4network string
	ipv6network string
	ipv4targets map[string][2]string
	ipv6targets map[string][2]string
	config      cfg.ConfigSettings
	events      publisher.Client
	done        chan struct{}
}

type Ping struct {
	target     string
	start_time time.Time
	rtt        time.Duration
}

type PingState struct {
	mu sync.RWMutex
	p  map[int]Ping
}

func New() *Pingbeat {
	return &Pingbeat{}
}

// Config reads in the pingbeat configuration file, validating
// configuration parameters and setting default values where needed
func (p *Pingbeat) Config(b *beat.Beat) error {
	// Read in provided config file, bail if problem
	err := cfgfile.Read(&p.config, "")
	if err != nil {
		logp.Err("Error reading configuration file: %v", err)
		return err
	}

	// Use period provided in config or default to 5s
	if p.config.Input.Period != nil {
		duration, err := time.ParseDuration(*p.config.Input.Period)
		p.period = duration
		if duration < time.Second || err != nil {
			logp.Warn("Config: Error parsing duration or duration too small: %v. Setting to default 10s", duration)
			p.period = 10 * time.Second
		}
	} else {
		logp.Warn("Config: No period set. Setting to default 10s")
		p.period = 10 * time.Second
	}
	logp.Debug("pingbeat", "Period %v\n", p.period)

	// Check if we can use privileged (i.e. raw socket) ping,
	// else use a UDP ping
	if *p.config.Input.Privileged {
		if os.Getuid() != 0 {
			err := errors.New("Privileged set but not running with privleges!")
			return err
		}
		p.ipv4network = "ip4:icmp"
		p.ipv6network = "ip6:ipv6-icmp"
	} else {
		p.ipv4network = "udp4"
		p.ipv6network = "udp6"

	}
	logp.Debug("pingbeat", "Using %v and/or %v for pings\n", p.ipv4network, p.ipv6network)

	// Check whether IPv4/IPv6 pings are requested in config
	// Default to just IPv4 pings
	if &p.config.Input.UseIPv4 != nil {
		p.useIPv4 = *p.config.Input.UseIPv4
	} else {
		p.useIPv4 = true
	}
	if &p.config.Input.UseIPv6 != nil {
		p.useIPv6 = *p.config.Input.UseIPv6
	} else {
		p.useIPv6 = false
	}
	logp.Debug("pingbeat", "Using IPv4: %v. Using IPv6: %v\n", p.useIPv4, p.useIPv6)

	// Fill the IPv4/IPv6 targets maps
	p.ipv4targets = make(map[string][2]string)
	p.ipv6targets = make(map[string][2]string)
	if p.config.Input.Targets != nil {
		for tag, targets := range *p.config.Input.Targets {
			for i := 0; i < len(targets); i++ {
				p.AddTarget(targets[i], tag)
			}
		}
	} else {
		err := errors.New("No targets specified, cannot continue!")
		return err
	}

	return nil
}

// Setup performs boilerplate Beats setup
func (p *Pingbeat) Setup(b *beat.Beat) error {
	p.events = b.Events
	p.done = make(chan struct{})
	return nil
}

func (p *Pingbeat) Run(b *beat.Beat) error {
	spool := pool.New()
	defer spool.Close()
	rpool := pool.New()
	defer rpool.Close()

	ticker := time.NewTicker(p.period)
	defer ticker.Stop()

	seq_no := 0

	createConn := func(n string, a string) *icmp.PacketConn {
		c, err := icmp.ListenPacket(n, a)
		if err != nil {
			logp.Err("Error creating connection: %v", err)
			return nil
		} else {
			return c
		}
	}

	c4 := &icmp.PacketConn{}
	if p.useIPv4 {
		c4 = createConn(p.ipv4network, "0.0.0.0")
	}
	defer c4.Close()

	c6 := &icmp.PacketConn{}
	if p.useIPv6 {
		c6 = createConn(p.ipv6network, "::")
	}
	defer c6.Close()

	pings := PingState{}
	pings.p = make(map[int]Ping)

	for {
		sendBatch := spool.Batch()
		select {
		case <-p.done:
			ticker.Stop()
			spool.Close()
			rpool.Close()
			return nil
		case <-ticker.C:

			if p.useIPv4 {
				go p.QueueRequests(&pings, ipv4.ICMPTypeEcho, &seq_no, c4, sendBatch)
			}
			if p.useIPv6 {
				go p.QueueRequests(&pings, ipv6.ICMPTypeEchoRequest, &seq_no, c6, sendBatch)
			}

			go func() {
				var r pool.WorkUnit
				for result := range sendBatch.Results() {
					if err := result.Error(); err != nil {
						logp.Err("Send unsuccessful: %v", err)
					} else {
						switch result.Value() {
						case ipv4.ICMPTypeEcho:
							r = rpool.Queue(RecvPing(c4))
						case ipv6.ICMPTypeEchoRequest:
							r = rpool.Queue(RecvPing(c6))
						default:
							logp.Err("Invalid ICMP message type")
						}
						r.Wait()
						if err := r.Error(); err != nil {
							logp.Err("Recv unsuccessful: %v", err)
						} else {
							reply := r.Value().(*PingReply)

							switch reply.text_payload.Body.(type) {
							case *icmp.TimeExceeded:
								logp.Err("time exceeded")
							case *icmp.PacketTooBig:
								logp.Err("packet too big")
							case *icmp.DstUnreach:
								logp.Err("unreachable")
							case *icmp.Echo:
								tgt := struct {
									s int
									t string
								}{reply.text_payload.Body.(*icmp.Echo).Seq, reply.target}
								go p.ProcessReplies(&pings, &tgt)
							default:
								logp.Err("Unknown packet response")
							}
						}
					}
				}
			}()
		}
	}
}

func (p *Pingbeat) Cleanup(b *beat.Beat) error {
	return nil
}

func (p *Pingbeat) Stop() {
	close(p.done)
}

// AddTarget takes a target name and tag, fetches the IP addresses associated
// with it and adds them to the Pingbeat struct
func (p *Pingbeat) AddTarget(target string, tag string) {
	if addr := net.ParseIP(target); addr.String() == target {
		if addr.To4() != nil && p.useIPv4 {
			p.ipv4targets[addr.String()] = [2]string{target, tag}
		} else if p.useIPv6 {
			p.ipv6targets[addr.String()] = [2]string{target, tag}
		}
	} else {
		ip4addr := make(chan string)
		ip6addr := make(chan string)
		go FetchIPs(ip4addr, ip6addr, target)
	lookup:
		for {
			select {
			case ip := <-ip4addr:
				if ip == "done" {
					break lookup
				} else if p.useIPv4 {
					logp.Debug("pingbeat", "Target %s has an IPv4 address %s\n", target, ip)
					p.ipv4targets[ip] = [2]string{target, tag}
				}
			case ip := <-ip6addr:
				if ip == "done" {
					break lookup
				} else if p.useIPv6 {
					logp.Debug("pingbeat", "Target %s has an IPv6 address %s\n", target, ip)
					p.ipv6targets[ip] = [2]string{target, tag}
				}
			}
		}
	}
}

// Addr2Name takes a address as a string and returns the name and tag
// associated with that address in the Pingbeat struct
func (p *Pingbeat) FetchDetails(addr string) (string, string) {
	var name, tag string
	if _, found := p.ipv4targets[addr]; found {
		name = p.ipv4targets[addr][0]
		tag = p.ipv4targets[addr][1]
	} else if _, found := p.ipv6targets[addr]; found {
		name = p.ipv6targets[addr][0]
		tag = p.ipv6targets[addr][1]
	} else {
		logp.Err("Error: %s not found in Pingbeat targets!", addr)
		name = "err"
		tag = "err"
	}
	return name, tag
}

// milliSeconds converts seconds to milliseconds
func milliSeconds(d time.Duration) float64 {
	msec := d / time.Millisecond
	nsec := d % time.Millisecond
	return float64(msec) + float64(nsec)*1e-6
}

// FetchIPs takes a target hostname, resolves the IP addresses for that
// hostname via DNS and returns the results through the ip4addr/ip6addr
// channels
func FetchIPs(ip4addr, ip6addr chan string, target string) {
	addrs, err := net.LookupIP(target)
	if err != nil {
		logp.Warn("Failed to resolve %s to IP address, ignoring this target.\n", target)
	} else {
		for j := 0; j < len(addrs); j++ {
			if addrs[j].To4() != nil {
				ip4addr <- addrs[j].String()
			} else {
				ip6addr <- addrs[j].String()
			}
		}
	}
	ip4addr <- "done"
	close(ip4addr)
	ip6addr <- "done"
	close(ip6addr)
	return
}

func (p *Pingbeat) QueueRequests(pings *PingState, ping_type icmp.Type, seq_no *int, conn *icmp.PacketConn, batch pool.Batch) {
	var network string
	targets := make(map[string][2]string)
	switch ping_type {
	case ipv4.ICMPTypeEcho:
		targets = p.ipv4targets
		network = p.ipv4network
	case ipv6.ICMPTypeEchoRequest:
		targets = p.ipv4targets
		network = p.ipv4network
	default:
		logp.Err("QueueTargets: Invalid ICMP message type")
	}
	for addr, _ := range targets {
		req, err := NewPingRequest(*seq_no, ping_type, addr, network)
		if err != nil {
			logp.Err("QueueTargets: %v", err)
		}
		batch.Queue(SendPing(conn, req))
		pings.mu.Lock()
		pings.p[*seq_no] = Ping{target: addr, start_time: time.Now().UTC()}
		pings.mu.Unlock()
		*seq_no++
		// reset sequence no if we go above a 32-bit value
		if *seq_no > 65535 {
			logp.Debug("pingbeat", "Resetting sequence number")
			*seq_no = 0
		}
	}
	batch.QueueComplete()
}

func (p *Pingbeat) ProcessReplies(state *PingState, tgt *struct {
	s int
	t string
}) {
	state.mu.Lock()
	rtt := time.Since(state.p[tgt.s].start_time)
	delete(state.p, tgt.s)
	state.mu.Unlock()
	if rtt > (5 * time.Second) {
		logp.Info("No record of ICMP packet: %i", tgt.s)
	}
	name, tag := p.FetchDetails(tgt.t)

	event := common.MapStr{
		"@timestamp":  common.Time(time.Now().UTC()),
		"type":        "pingbeat",
		"target_name": name,
		"target_addr": tgt.t,
		"tag":         tag,
		"rtt":         milliSeconds(rtt),
	}
	p.events.PublishEvent(event)
}

func SendPing(c *icmp.PacketConn, req *PingRequest) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		if _, err := c.WriteTo(req.binary_payload, req.addr); err != nil {
			return nil, err
		}
		if err := c.SetReadDeadline(time.Now().Add(5000 * time.Millisecond)); err != nil {
			return nil, err
		}
		if wu.IsCancelled() {
			logp.Debug("pingbeat", "SendPing: workunit cancelled")
			return nil, nil
		}
		return req.ping_type, nil
	}
}

func RecvPing(c *icmp.PacketConn) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		var ping_type icmp.Type
		switch {
		case c.IPv4PacketConn() != nil:
			ping_type = ipv4.ICMPTypeEcho
		case c.IPv4PacketConn() != nil:
			ping_type = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("RecvPing: Unknown connection type")
			return nil, err
		}
		bytes := make([]byte, 1500)
		n, peer, err := c.ReadFrom(bytes)
		if err != nil {
			return nil, err
		}
		rep, err := NewPingReply(n, peer.String(), bytes, ping_type)
		if err != nil {
			return nil, err
		}
		if wu.IsCancelled() {
			logp.Debug("pingbeat", "RecvPing: workunit cancelled")
			return nil, nil
		}
		bytes = nil
		return rep, nil
	}
}
<commit_msg>Rename some variables.  Handle missing ping packets.<commit_after>package pingbeat

import (
	"errors"
	"github.com/davecgh/go-spew/spew"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/cfgfile"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	cfg "github.com/joshuar/pingbeat/config"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
	"gopkg.in/go-playground/pool.v3"
	"net"
	"os"
	"sync"
	"time"
)

// Pingbeat struct contains all options and
// hosts to ping
type Pingbeat struct {
	isAlive     bool
	useIPv4     bool
	useIPv6     bool
	period      time.Duration
	ipv4network string
	ipv6network string
	ipv4targets map[string][2]string
	ipv6targets map[string][2]string
	config      cfg.ConfigSettings
	events      publisher.Client
	done        chan struct{}
}

type Ping struct {
	target     string
	start_time time.Time
	rtt        time.Duration
}

type PingState struct {
	mu sync.RWMutex
	p  map[int]Ping
}

func New() *Pingbeat {
	return &Pingbeat{}
}

// Config reads in the pingbeat configuration file, validating
// configuration parameters and setting default values where needed
func (p *Pingbeat) Config(b *beat.Beat) error {
	// Read in provided config file, bail if problem
	err := cfgfile.Read(&p.config, "")
	if err != nil {
		logp.Err("Error reading configuration file: %v", err)
		return err
	}

	// Use period provided in config or default to 10s
	if p.config.Input.Period != nil {
		duration, err := time.ParseDuration(*p.config.Input.Period)
		p.period = duration
		if duration < time.Second || err != nil {
			logp.Warn("Config: Error parsing period or period too small: %v. Setting to default 10s", duration)
			p.period = 10 * time.Second
		}
	} else {
		logp.Warn("Config: No period set. Setting to default 10s")
		p.period = 10 * time.Second
	}
	logp.Debug("pingbeat", "Period %v\n", p.period)

	// Check if we can use privileged (i.e. raw socket) ping,
	// else use a UDP ping
	if *p.config.Input.Privileged {
		if os.Getuid() != 0 {
			err := errors.New("Privileged set but not running with privleges!")
			return err
		}
		p.ipv4network = "ip4:icmp"
		p.ipv6network = "ip6:ipv6-icmp"
	} else {
		p.ipv4network = "udp4"
		p.ipv6network = "udp6"

	}
	logp.Debug("pingbeat", "Using %v and/or %v for pings\n", p.ipv4network, p.ipv6network)

	// Check whether IPv4/IPv6 pings are requested in config
	// Default to just IPv4 pings
	if &p.config.Input.UseIPv4 != nil {
		p.useIPv4 = *p.config.Input.UseIPv4
	} else {
		p.useIPv4 = true
	}
	if &p.config.Input.UseIPv6 != nil {
		p.useIPv6 = *p.config.Input.UseIPv6
	} else {
		p.useIPv6 = false
	}
	logp.Debug("pingbeat", "Using IPv4: %v. Using IPv6: %v\n", p.useIPv4, p.useIPv6)

	// Fill the IPv4/IPv6 targets maps
	p.ipv4targets = make(map[string][2]string)
	p.ipv6targets = make(map[string][2]string)
	if p.config.Input.Targets != nil {
		for tag, targets := range *p.config.Input.Targets {
			for i := 0; i < len(targets); i++ {
				p.AddTarget(targets[i], tag)
			}
		}
	} else {
		err := errors.New("No targets specified, cannot continue!")
		return err
	}

	return nil
}

// Setup performs boilerplate Beats setup
func (p *Pingbeat) Setup(b *beat.Beat) error {
	p.events = b.Events
	p.done = make(chan struct{})
	return nil
}

func (p *Pingbeat) Run(b *beat.Beat) error {
	spool := pool.New()
	defer spool.Close()
	rpool := pool.New()
	defer rpool.Close()

	ticker := time.NewTicker(p.period)
	defer ticker.Stop()

	seq_no := 0

	createConn := func(n string, a string) *icmp.PacketConn {
		c, err := icmp.ListenPacket(n, a)
		if err != nil {
			logp.Err("Error creating connection: %v", err)
			return nil
		} else {
			return c
		}
	}

	c4 := &icmp.PacketConn{}
	if p.useIPv4 {
		c4 = createConn(p.ipv4network, "0.0.0.0")
	}
	defer c4.Close()

	c6 := &icmp.PacketConn{}
	if p.useIPv6 {
		c6 = createConn(p.ipv6network, "::")
	}
	defer c6.Close()

	pings := PingState{}
	pings.p = make(map[int]Ping)

	for {
		select {
		case <-p.done:
			ticker.Stop()
			spool.Close()
			rpool.Close()
			return nil
		case <-ticker.C:
			sendBatch := spool.Batch()

			if p.useIPv4 {
				go p.QueueRequests(&pings, &seq_no, c4, sendBatch)
			}
			if p.useIPv6 {
				go p.QueueRequests(&pings, &seq_no, c6, sendBatch)
			}

			var r pool.WorkUnit
			for result := range sendBatch.Results() {
				if err := result.Error(); err != nil {
					logp.Err("Send unsuccessful: %v", err)
				} else {
					switch result.Value() {
					case ipv4.ICMPTypeEcho:
						r = rpool.Queue(RecvPing(c4))
					case ipv6.ICMPTypeEchoRequest:
						r = rpool.Queue(RecvPing(c6))
					default:
						logp.Err("Invalid ICMP message type")
					}
					r.Wait()
					if err := r.Error(); err != nil {
						logp.Err("Recv unsuccessful: %v", err)
					} else {
						reply := r.Value().(*PingReply)
						switch reply.text_payload.Body.(type) {
						case *icmp.TimeExceeded:
							logp.Err("time exceeded")
						case *icmp.PacketTooBig:
							logp.Err("packet too big")
						case *icmp.DstUnreach:
							logp.Err("unreachable")
							spew.Dump(reply)
						case *icmp.Echo:
							tgt := struct {
								s int
								t string
							}{reply.text_payload.Body.(*icmp.Echo).Seq, reply.target}
							go p.ProcessPing(&pings, &tgt)
						default:
							logp.Err("Unknown packet response")
						}
					}
					r.Cancel()
				}
			}
			go p.ProcessMissing(&pings)
		}
	}
}

func (p *Pingbeat) Cleanup(b *beat.Beat) error {
	return nil
}

func (p *Pingbeat) Stop() {
	close(p.done)
}

// AddTarget takes a target name and tag, fetches the IP addresses associated
// with it and adds them to the Pingbeat struct
func (p *Pingbeat) AddTarget(target string, tag string) {
	if addr := net.ParseIP(target); addr.String() == target {
		if addr.To4() != nil && p.useIPv4 {
			p.ipv4targets[addr.String()] = [2]string{target, tag}
		} else if p.useIPv6 {
			p.ipv6targets[addr.String()] = [2]string{target, tag}
		}
	} else {
		ip4addr := make(chan string)
		ip6addr := make(chan string)
		go FetchIPs(ip4addr, ip6addr, target)
	lookup:
		for {
			select {
			case ip := <-ip4addr:
				if ip == "done" {
					break lookup
				} else if p.useIPv4 {
					logp.Debug("pingbeat", "Target %s has an IPv4 address %s\n", target, ip)
					p.ipv4targets[ip] = [2]string{target, tag}
				}
			case ip := <-ip6addr:
				if ip == "done" {
					break lookup
				} else if p.useIPv6 {
					logp.Debug("pingbeat", "Target %s has an IPv6 address %s\n", target, ip)
					p.ipv6targets[ip] = [2]string{target, tag}
				}
			}
		}
	}
}

// Addr2Name takes a address as a string and returns the name and tag
// associated with that address in the Pingbeat struct
func (p *Pingbeat) FetchDetails(addr string) (string, string) {
	var name, tag string
	if _, found := p.ipv4targets[addr]; found {
		name = p.ipv4targets[addr][0]
		tag = p.ipv4targets[addr][1]
	} else if _, found := p.ipv6targets[addr]; found {
		name = p.ipv6targets[addr][0]
		tag = p.ipv6targets[addr][1]
	} else {
		logp.Err("Error: %s not found in Pingbeat targets!", addr)
		name = "err"
		tag = "err"
	}
	return name, tag
}

// milliSeconds converts seconds to milliseconds
func milliSeconds(d time.Duration) float64 {
	msec := d / time.Millisecond
	nsec := d % time.Millisecond
	return float64(msec) + float64(nsec)*1e-6
}

// FetchIPs takes a target hostname, resolves the IP addresses for that
// hostname via DNS and returns the results through the ip4addr/ip6addr
// channels
func FetchIPs(ip4addr, ip6addr chan string, target string) {
	addrs, err := net.LookupIP(target)
	if err != nil {
		logp.Warn("Failed to resolve %s to IP address, ignoring this target.\n", target)
	} else {
		for j := 0; j < len(addrs); j++ {
			if addrs[j].To4() != nil {
				ip4addr <- addrs[j].String()
			} else {
				ip6addr <- addrs[j].String()
			}
		}
	}
	ip4addr <- "done"
	close(ip4addr)
	ip6addr <- "done"
	close(ip6addr)
	return
}

func (p *Pingbeat) QueueRequests(pings *PingState, seq_no *int, conn *icmp.PacketConn, batch pool.Batch) {
	var network string
	var ping_type icmp.Type
	switch {
	case conn.IPv4PacketConn() != nil:
		ping_type = ipv4.ICMPTypeEcho
	case conn.IPv4PacketConn() != nil:
		ping_type = ipv6.ICMPTypeEchoRequest
	default:
		logp.Err("QueueRequests: Unknown connection type")
	}
	targets := make(map[string][2]string)
	switch ping_type {
	case ipv4.ICMPTypeEcho:
		targets = p.ipv4targets
		network = p.ipv4network
	case ipv6.ICMPTypeEchoRequest:
		targets = p.ipv6targets
		network = p.ipv6network
	default:
		logp.Err("QueueTargets: Invalid ICMP message type")
	}
	for addr, _ := range targets {
		req, err := NewPingRequest(*seq_no, ping_type, addr, network)
		if err != nil {
			logp.Err("QueueTargets: %v", err)
		}
		batch.Queue(SendPing(conn, p.period, req))
		pings.mu.Lock()
		pings.p[*seq_no] = Ping{target: addr, start_time: time.Now().UTC()}
		pings.mu.Unlock()
		*seq_no++
		// reset sequence no if we go above a 32-bit value
		if *seq_no > 65535 {
			logp.Debug("pingbeat", "Resetting sequence number")
			*seq_no = 0
		}
	}
	batch.QueueComplete()
}

func (p *Pingbeat) ProcessPing(state *PingState, tgt *struct {
	s int
	t string
}) {
	state.mu.Lock()
	rtt := time.Since(state.p[tgt.s].start_time)
	delete(state.p, tgt.s)
	state.mu.Unlock()
	if rtt > (5 * time.Second) {
		logp.Info("No record of ICMP packet: %i", tgt.s)
	}
	name, tag := p.FetchDetails(tgt.t)

	event := common.MapStr{
		"@timestamp":  common.Time(time.Now().UTC()),
		"type":        "pingbeat",
		"target_name": name,
		"target_addr": tgt.t,
		"tag":         tag,
		"rtt":         milliSeconds(rtt),
	}
	p.events.PublishEvent(event)
}

func (p *Pingbeat) ProcessMissing(state *PingState) {
	for seq_no, info := range state.p {
		name, tag := p.FetchDetails(info.target)
		event := common.MapStr{
			"@timestamp":  common.Time(time.Now().UTC()),
			"type":        "pingbeat",
			"target_name": name,
			"target_addr": info.target,
			"tag":         tag,
			"loss":        true,
		}
		p.events.PublishEvent(event)
		state.mu.Lock()
		delete(state.p, seq_no)
		state.mu.Unlock()
	}
}

func SendPing(conn *icmp.PacketConn, timeout time.Duration, req *PingRequest) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		if _, err := conn.WriteTo(req.binary_payload, req.addr); err != nil {
			return nil, err
		}
		if err := conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
			return nil, err
		}
		if wu.IsCancelled() {
			logp.Debug("pingbeat", "SendPing: workunit cancelled")
			return nil, nil
		}
		return req.ping_type, nil
	}
}

func RecvPing(conn *icmp.PacketConn) pool.WorkFunc {
	return func(wu pool.WorkUnit) (interface{}, error) {
		var ping_type icmp.Type
		switch {
		case conn.IPv4PacketConn() != nil:
			ping_type = ipv4.ICMPTypeEcho
		case conn.IPv4PacketConn() != nil:
			ping_type = ipv6.ICMPTypeEchoRequest
		default:
			err := errors.New("RecvPing: Unknown connection type")
			return nil, err
		}
		bytes := make([]byte, 1500)
		n, peer, err := conn.ReadFrom(bytes)
		if err != nil {
			return nil, err
		}
		rep, err := NewPingReply(n, peer.String(), bytes, ping_type)
		if err != nil {
			return nil, err
		}
		if wu.IsCancelled() {
			logp.Debug("pingbeat", "RecvPing: workunit cancelled")
			return nil, nil
		}
		bytes = nil
		return rep, nil
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Skaffold 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 portforward

import (
	"context"
	"fmt"
	"io/ioutil"
	"sync"
	"testing"
	"time"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/deploy"
	"github.com/google/go-cmp/cmp"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
	fakekubeclientset "k8s.io/client-go/kubernetes/fake"
)

type testForwarder struct {
	forwardedResources forwardedResources
	forwardedPorts     forwardedPorts
}

func (f *testForwarder) Forward(ctx context.Context, pfe *portForwardEntry) {
	f.forwardedResources.Store(pfe.key(), pfe)
	f.forwardedPorts.Store(pfe.localPort, true)
}

func (f *testForwarder) Monitor(_ *portForwardEntry, _ func()) {}

func (f *testForwarder) Terminate(pfe *portForwardEntry) {
	f.forwardedResources.Delete(pfe.key())
	f.forwardedPorts.Delete(pfe.resource.Port)
}

func newTestForwarder() *testForwarder {
	return &testForwarder{
		forwardedResources: newForwardedResources(),
		forwardedPorts:     newForwardedPorts(),
	}
}

func mockRetrieveAvailablePort(taken map[int]struct{}, availablePorts []int) func(int, util.ForwardedPorts) int {
	// Return first available port in ports that isn't taken
	lock := sync.Mutex{}
	return func(int, util.ForwardedPorts) int {
		for _, p := range availablePorts {
			lock.Lock()
			if _, ok := taken[p]; ok {
				lock.Unlock()
				continue
			}
			taken[p] = struct{}{}
			lock.Unlock()
			return p
		}
		return -1
	}
}

func TestStart(t *testing.T) {
	svc1 := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc1",
		Namespace: "default",
		Port:      8080,
	}

	svc2 := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc2",
		Namespace: "default",
		Port:      9000,
	}

	tests := []struct {
		description    string
		resources      []*latest.PortForwardResource
		availablePorts []int
		expected       map[string]*portForwardEntry
	}{
		{
			description:    "forward two services",
			resources:      []*latest.PortForwardResource{svc1, svc2},
			availablePorts: []int{8080, 9000},
			expected: map[string]*portForwardEntry{
				"service-svc1-default-8080": {
					resource:  *svc1,
					localPort: 8080,
				},
				"service-svc2-default-9000": {
					resource:  *svc2,
					localPort: 9000,
				},
			},
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			event.InitializeState(latest.BuildConfig{})
			fakeForwarder := newTestForwarder()
			rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", nil)
			rf.EntryForwarder = fakeForwarder

			t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, test.availablePorts))
			t.Override(&retrieveServices, func(string, []string) ([]*latest.PortForwardResource, error) {
				return test.resources, nil
			})

			if err := rf.Start(context.Background()); err != nil {
				t.Fatalf("error starting resource forwarder: %v", err)
			}
			// poll up to 10 seconds for the resources to be forwarded
			err := wait.PollImmediate(100*time.Millisecond, 10*time.Second, func() (bool, error) {
				return len(test.expected) == fakeForwarder.forwardedResources.Length(), nil
			})
			if err != nil {
				t.Fatalf("expected entries didn't match actual entries. Expected: \n %v Actual: \n %v", test.expected, fakeForwarder.forwardedResources)
			}
		})
	}
}

func TestGetCurrentEntryFunc(t *testing.T) {
	tests := []struct {
		description        string
		forwardedResources map[string]*portForwardEntry
		availablePorts     []int
		resource           latest.PortForwardResource
		expected           *portForwardEntry
	}{
		{
			description: "port forward service",
			resource: latest.PortForwardResource{
				Type: "service",
				Name: "serviceName",
				Port: 8080,
			},
			availablePorts: []int{8080},
			expected: &portForwardEntry{
				localPort:       8080,
				terminationLock: &sync.Mutex{},
			},
		}, {
			description: "port forward existing deployment",
			resource: latest.PortForwardResource{
				Type:      "deployment",
				Namespace: "default",
				Name:      "depName",
				Port:      8080,
			},
			forwardedResources: map[string]*portForwardEntry{
				"deployment-depName-default-8080": {
					resource: latest.PortForwardResource{
						Type:      "deployment",
						Namespace: "default",
						Name:      "depName",
						Port:      8080,
					},
					localPort: 9000,
				},
			},
			expected: &portForwardEntry{
				localPort:       9000,
				terminationLock: &sync.Mutex{},
			},
		},
	}

	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			expectedEntry := test.expected
			expectedEntry.resource = test.resource

			rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", nil)
			rf.forwardedResources = forwardedResources{
				resources: test.forwardedResources,
				lock:      &sync.Mutex{},
			}

			t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, test.availablePorts))

			actualEntry := rf.getCurrentEntry(test.resource)
			t.CheckDeepEqual(expectedEntry, actualEntry, cmp.AllowUnexported(portForwardEntry{}, sync.Mutex{}))
		})
	}
}

func TestUserDefinedResources(t *testing.T) {
	svc := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc1",
		Namespace: "default",
		Port:      8080,
	}

	pod := &latest.PortForwardResource{
		Type:      constants.Pod,
		Name:      "pod",
		Namespace: "default",
		Port:      9000,
	}

	expected := map[string]*portForwardEntry{
		"service-svc1-default-8080": {
			resource:  *svc,
			localPort: 8080,
		},
		"pod-pod-default-9000": {
			resource:  *pod,
			localPort: 9000,
		},
	}

	testutil.Run(t, "one service and one user defined pod", func(t *testutil.T) {
		event.InitializeState(latest.BuildConfig{})
		fakeForwarder := newTestForwarder()
		rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", []*latest.PortForwardResource{pod})
		rf.EntryForwarder = fakeForwarder

		t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, []int{8080, 9000}))
		t.Override(&retrieveServices, func(string, []string) ([]*latest.PortForwardResource, error) {
			return []*latest.PortForwardResource{svc}, nil
		})

		if err := rf.Start(context.Background()); err != nil {
			t.Fatalf("error starting resource forwarder: %v", err)
		}
		// poll up to 10 seconds for the resources to be forwarded
		err := wait.PollImmediate(100*time.Millisecond, 10*time.Second, func() (bool, error) {
			return len(expected) == fakeForwarder.forwardedResources.Length(), nil
		})
		if err != nil {
			t.Fatalf("expected entries didn't match actual entries. Expected: \n %v Actual: \n %v", expected, fakeForwarder.forwardedResources.resources)
		}
	})
}

func mockClient(m kubernetes.Interface) func() (kubernetes.Interface, error) {
	return func() (kubernetes.Interface, error) {
		return m, nil
	}

}

func TestRetrieveServices(t *testing.T) {
	tests := []struct {
		description string
		namespaces  []string
		services    []*v1.Service
		expected    []*latest.PortForwardResource
	}{
		{
			description: "multiple services in multiple namespaces",
			namespaces:  []string{"test", "test1"},
			services: []*v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc1",
						Namespace: "test",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8080}}},
				}, {
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc2",
						Namespace: "test1",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8081}}},
				},
			},
			expected: []*latest.PortForwardResource{{
				Type:      constants.Service,
				Name:      "svc1",
				Namespace: "test",
				Port:      8080,
				LocalPort: 8080,
			}, {
				Type:      constants.Service,
				Name:      "svc2",
				Namespace: "test1",
				Port:      8081,
				LocalPort: 8081,
			}},
		}, {
			description: "no services in given namespace",
			namespaces:  []string{"randon"},
			services: []*v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc1",
						Namespace: "test",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8080}}},
				},
			},
		},
	}

	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			objs := make([]runtime.Object, len(test.services))
			for i, s := range test.services {
				objs[i] = s
			}
			client := fakekubeclientset.NewSimpleClientset(objs...)
			t.Override(&kClientSet, mockClient(client))
			actual, err := retrieveServiceResources(fmt.Sprintf("%s=9876-6789", deploy.RunIDLabel), test.namespaces)
			t.CheckNoError(err)
			t.CheckDeepEqual(test.expected, actual)
		})
	}
}
<commit_msg>add another test case for services with no ports exposeD<commit_after>/*
Copyright 2019 The Skaffold 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 portforward

import (
	"context"
	"fmt"
	"io/ioutil"
	"sync"
	"testing"
	"time"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/deploy"
	"github.com/google/go-cmp/cmp"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/event"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
	"github.com/GoogleContainerTools/skaffold/testutil"
	fakekubeclientset "k8s.io/client-go/kubernetes/fake"
)

type testForwarder struct {
	forwardedResources forwardedResources
	forwardedPorts     forwardedPorts
}

func (f *testForwarder) Forward(ctx context.Context, pfe *portForwardEntry) {
	f.forwardedResources.Store(pfe.key(), pfe)
	f.forwardedPorts.Store(pfe.localPort, true)
}

func (f *testForwarder) Monitor(_ *portForwardEntry, _ func()) {}

func (f *testForwarder) Terminate(pfe *portForwardEntry) {
	f.forwardedResources.Delete(pfe.key())
	f.forwardedPorts.Delete(pfe.resource.Port)
}

func newTestForwarder() *testForwarder {
	return &testForwarder{
		forwardedResources: newForwardedResources(),
		forwardedPorts:     newForwardedPorts(),
	}
}

func mockRetrieveAvailablePort(taken map[int]struct{}, availablePorts []int) func(int, util.ForwardedPorts) int {
	// Return first available port in ports that isn't taken
	lock := sync.Mutex{}
	return func(int, util.ForwardedPorts) int {
		for _, p := range availablePorts {
			lock.Lock()
			if _, ok := taken[p]; ok {
				lock.Unlock()
				continue
			}
			taken[p] = struct{}{}
			lock.Unlock()
			return p
		}
		return -1
	}
}

func TestStart(t *testing.T) {
	svc1 := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc1",
		Namespace: "default",
		Port:      8080,
	}

	svc2 := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc2",
		Namespace: "default",
		Port:      9000,
	}

	tests := []struct {
		description    string
		resources      []*latest.PortForwardResource
		availablePorts []int
		expected       map[string]*portForwardEntry
	}{
		{
			description:    "forward two services",
			resources:      []*latest.PortForwardResource{svc1, svc2},
			availablePorts: []int{8080, 9000},
			expected: map[string]*portForwardEntry{
				"service-svc1-default-8080": {
					resource:  *svc1,
					localPort: 8080,
				},
				"service-svc2-default-9000": {
					resource:  *svc2,
					localPort: 9000,
				},
			},
		},
	}
	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			event.InitializeState(latest.BuildConfig{})
			fakeForwarder := newTestForwarder()
			rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", nil)
			rf.EntryForwarder = fakeForwarder

			t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, test.availablePorts))
			t.Override(&retrieveServices, func(string, []string) ([]*latest.PortForwardResource, error) {
				return test.resources, nil
			})

			if err := rf.Start(context.Background()); err != nil {
				t.Fatalf("error starting resource forwarder: %v", err)
			}
			// poll up to 10 seconds for the resources to be forwarded
			err := wait.PollImmediate(100*time.Millisecond, 10*time.Second, func() (bool, error) {
				return len(test.expected) == fakeForwarder.forwardedResources.Length(), nil
			})
			if err != nil {
				t.Fatalf("expected entries didn't match actual entries. Expected: \n %v Actual: \n %v", test.expected, fakeForwarder.forwardedResources)
			}
		})
	}
}

func TestGetCurrentEntryFunc(t *testing.T) {
	tests := []struct {
		description        string
		forwardedResources map[string]*portForwardEntry
		availablePorts     []int
		resource           latest.PortForwardResource
		expected           *portForwardEntry
	}{
		{
			description: "port forward service",
			resource: latest.PortForwardResource{
				Type: "service",
				Name: "serviceName",
				Port: 8080,
			},
			availablePorts: []int{8080},
			expected: &portForwardEntry{
				localPort:       8080,
				terminationLock: &sync.Mutex{},
			},
		}, {
			description: "port forward existing deployment",
			resource: latest.PortForwardResource{
				Type:      "deployment",
				Namespace: "default",
				Name:      "depName",
				Port:      8080,
			},
			forwardedResources: map[string]*portForwardEntry{
				"deployment-depName-default-8080": {
					resource: latest.PortForwardResource{
						Type:      "deployment",
						Namespace: "default",
						Name:      "depName",
						Port:      8080,
					},
					localPort: 9000,
				},
			},
			expected: &portForwardEntry{
				localPort:       9000,
				terminationLock: &sync.Mutex{},
			},
		},
	}

	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			expectedEntry := test.expected
			expectedEntry.resource = test.resource

			rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", nil)
			rf.forwardedResources = forwardedResources{
				resources: test.forwardedResources,
				lock:      &sync.Mutex{},
			}

			t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, test.availablePorts))

			actualEntry := rf.getCurrentEntry(test.resource)
			t.CheckDeepEqual(expectedEntry, actualEntry, cmp.AllowUnexported(portForwardEntry{}, sync.Mutex{}))
		})
	}
}

func TestUserDefinedResources(t *testing.T) {
	svc := &latest.PortForwardResource{
		Type:      constants.Service,
		Name:      "svc1",
		Namespace: "default",
		Port:      8080,
	}

	pod := &latest.PortForwardResource{
		Type:      constants.Pod,
		Name:      "pod",
		Namespace: "default",
		Port:      9000,
	}

	expected := map[string]*portForwardEntry{
		"service-svc1-default-8080": {
			resource:  *svc,
			localPort: 8080,
		},
		"pod-pod-default-9000": {
			resource:  *pod,
			localPort: 9000,
		},
	}

	testutil.Run(t, "one service and one user defined pod", func(t *testutil.T) {
		event.InitializeState(latest.BuildConfig{})
		fakeForwarder := newTestForwarder()
		rf := NewResourceForwarder(NewEntryManager(ioutil.Discard, nil), []string{"test"}, "", []*latest.PortForwardResource{pod})
		rf.EntryForwarder = fakeForwarder

		t.Override(&retrieveAvailablePort, mockRetrieveAvailablePort(map[int]struct{}{}, []int{8080, 9000}))
		t.Override(&retrieveServices, func(string, []string) ([]*latest.PortForwardResource, error) {
			return []*latest.PortForwardResource{svc}, nil
		})

		if err := rf.Start(context.Background()); err != nil {
			t.Fatalf("error starting resource forwarder: %v", err)
		}
		// poll up to 10 seconds for the resources to be forwarded
		err := wait.PollImmediate(100*time.Millisecond, 10*time.Second, func() (bool, error) {
			return len(expected) == fakeForwarder.forwardedResources.Length(), nil
		})
		if err != nil {
			t.Fatalf("expected entries didn't match actual entries. Expected: \n %v Actual: \n %v", expected, fakeForwarder.forwardedResources.resources)
		}
	})
}

func mockClient(m kubernetes.Interface) func() (kubernetes.Interface, error) {
	return func() (kubernetes.Interface, error) {
		return m, nil
	}

}

func TestRetrieveServices(t *testing.T) {
	tests := []struct {
		description string
		namespaces  []string
		services    []*v1.Service
		expected    []*latest.PortForwardResource
	}{
		{
			description: "multiple services in multiple namespaces",
			namespaces:  []string{"test", "test1"},
			services: []*v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc1",
						Namespace: "test",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8080}}},
				}, {
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc2",
						Namespace: "test1",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8081}}},
				},
			},
			expected: []*latest.PortForwardResource{{
				Type:      constants.Service,
				Name:      "svc1",
				Namespace: "test",
				Port:      8080,
				LocalPort: 8080,
			}, {
				Type:      constants.Service,
				Name:      "svc2",
				Namespace: "test1",
				Port:      8081,
				LocalPort: 8081,
			}},
		}, {
			description: "no services in given namespace",
			namespaces:  []string{"randon"},
			services: []*v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc1",
						Namespace: "test",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
					Spec: v1.ServiceSpec{Ports: []v1.ServicePort{{Port: 8080}}},
				},
			},
		}, {
			description: "services present but does not expose any port",
			namespaces:  []string{"test"},
			services: []*v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "svc1",
						Namespace: "test",
						Labels: map[string]string{
							deploy.RunIDLabel: "9876-6789",
						},
					},
				},
			},
		},
	}

	for _, test := range tests {
		testutil.Run(t, test.description, func(t *testutil.T) {
			objs := make([]runtime.Object, len(test.services))
			for i, s := range test.services {
				objs[i] = s
			}
			client := fakekubeclientset.NewSimpleClientset(objs...)
			t.Override(&kClientSet, mockClient(client))
			actual, err := retrieveServiceResources(fmt.Sprintf("%s=9876-6789", deploy.RunIDLabel), test.namespaces)
			t.CheckNoError(err)
			t.CheckDeepEqual(test.expected, actual)
		})
	}
}
<|endoftext|>


<commit_before>//+build !debug

package debug

const Enabled = false

type guard struct{}
func (g guard) IRelease(f string, args ...interface{}) {}

// IPrintf is no op unless you comple with the `debug` tag
func IPrintf(f string, args ...interface{}) guard { return nil }

// Printf is no op unless you compile with the `debug` tag
func Printf(f string, args ...interface{}) {}

// Dump dumps the objects using go-spew
func Dump(v ...interface{}) {}
<commit_msg>appease the compiler<commit_after>//+build !debug

package debug

const Enabled = false

type guard struct{}
func (g guard) IRelease(f string, args ...interface{}) {}

// IPrintf is no op unless you comple with the `debug` tag
func IPrintf(f string, args ...interface{}) guard { return guard{} }

// Printf is no op unless you compile with the `debug` tag
func Printf(f string, args ...interface{}) {}

// Dump dumps the objects using go-spew
func Dump(v ...interface{}) {}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 main

import (
	"flag"
	"os"

	log "github.com/sirupsen/logrus"

	secretsstore "github.com/deislabs/secrets-store-csi-driver/pkg/secrets-store"
)

var (
	endpoint           = flag.String("endpoint", "unix://tmp/csi.sock", "CSI endpoint")
	driverName         = flag.String("drivername", "secrets-store.csi.k8s.com", "name of the driver")
	nodeID             = flag.String("nodeid", "", "node id")
	debug              = flag.Bool("debug", false, "sets log to debug level")
	logFormatJSON      = flag.Bool("log-format-json", false, "set log formatter to json")
	logReportCaller    = flag.Bool("log-report-caller", false, "include the calling method as fields in the log")
	providerVolumePath = flag.String("provider-volume", "/etc/kubernetes/secrets-store-csi-providers", "Volume path for provider")
)

func main() {
	flag.Parse()

	log.SetLevel(log.InfoLevel)
	if *debug {
		log.SetLevel(log.DebugLevel)
	}
	if *logFormatJSON {
		log.SetFormatter(&log.JSONFormatter{})
	}

	log.SetReportCaller(*logReportCaller)

	handle()
	os.Exit(0)
}

func handle() {
	driver := secretsstore.GetDriver()
	driver.Run(*driverName, *nodeID, *endpoint, *providerVolumePath)
}
<commit_msg>remove os.exit<commit_after>/*
Copyright 2018 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 main

import (
	"flag"

	log "github.com/sirupsen/logrus"

	secretsstore "github.com/deislabs/secrets-store-csi-driver/pkg/secrets-store"
)

var (
	endpoint           = flag.String("endpoint", "unix://tmp/csi.sock", "CSI endpoint")
	driverName         = flag.String("drivername", "secrets-store.csi.k8s.com", "name of the driver")
	nodeID             = flag.String("nodeid", "", "node id")
	debug              = flag.Bool("debug", false, "sets log to debug level")
	logFormatJSON      = flag.Bool("log-format-json", false, "set log formatter to json")
	logReportCaller    = flag.Bool("log-report-caller", false, "include the calling method as fields in the log")
	providerVolumePath = flag.String("provider-volume", "/etc/kubernetes/secrets-store-csi-providers", "Volume path for provider")
)

func main() {
	flag.Parse()

	log.SetLevel(log.InfoLevel)
	if *debug {
		log.SetLevel(log.DebugLevel)
	}
	if *logFormatJSON {
		log.SetFormatter(&log.JSONFormatter{})
	}

	log.SetReportCaller(*logReportCaller)

	handle()
}

func handle() {
	driver := secretsstore.GetDriver()
	driver.Run(*driverName, *nodeID, *endpoint, *providerVolumePath)
}
<|endoftext|>