<commit_before>// Copyright (c) 2020-2021 Tigera, Inc. All rights reserved.

// 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 state

import (
	"unsafe"

	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/felix/bpf"
)

type PolicyResult int32

const (
	PolicyNoMatch PolicyResult = iota
	PolicyAllow
	PolicyDeny
	PolicyTailCallFailed = 10
)

// struct cali_tc_state {
//    __be32 ip_src;4
//    __be32 ip_dst;8
//    __be32 post_nat_ip_dst;12
//    __be32 nat_tun_src;16
//    enum calico_policy_result pol_rc;24
//    __u16 sport;26
//    __u16 dport;28
//    __u16 post_nat_dport;30
//    __u8 ip_proto;31
//    __u8 pad;
//    __be16 ip_size;
//    struct calico_ct_result ct_result;
//    struct calico_nat_dest nat_dest;
//    __u64 prog_start_time;
// };
type State struct {
	SrcAddr             uint32
	DstAddr             uint32
	PreNATDstAddr       uint32
	PostNATDstAddr      uint32
	NATTunSrcAddr       uint32
	PolicyRC            PolicyResult
	SrcPort             uint16
	DstPort             uint16
	PreNATDstPort       uint16
	PostNATDstPort      uint16
	IPProto             uint8
	Flags               uint8
	IPSize              uint16
	ConntrackRCFlags    uint32
	ConntrackNATIPPort  uint64
	ConntrackTunIP      uint32
	ConntrackIfIndexFwd uint32
	ConntrackIfIndexCtd uint32
	NATData             uint64
	ProgStartTime       uint64
}

const expectedSize = 80

func (s *State) AsBytes() []byte {
	size := unsafe.Sizeof(State{})
	if size != expectedSize {
		log.WithField("size", size).Panic("Incorrect struct size")
	}
	bPtr := (*[expectedSize]byte)(unsafe.Pointer(s))
	bytes := make([]byte, expectedSize)
	copy(bytes, bPtr[:])
	return bytes
}

func StateFromBytes(bytes []byte) State {
	s := State{}
	bPtr := (*[expectedSize]byte)(unsafe.Pointer(&s))
	copy(bPtr[:], bytes)
	return s
}

func Map(mc *bpf.MapContext) bpf.Map {
	return mc.NewPinnedMap(bpf.MapParameters{
		Filename:   "/sys/fs/bpf/tc/globals/cali_v4_state",
		Type:       "percpu_array",
		KeySize:    4,
		ValueSize:  expectedSize,
		MaxEntries: 1,
		Name:       "cali_v4_state",
		Version:    3,
	})
}

func MapForTest(mc *bpf.MapContext) bpf.Map {
	return mc.NewPinnedMap(bpf.MapParameters{
		Filename:   "/sys/fs/bpf/tc/globals/test_v4_state",
		Type:       "array",
		KeySize:    4,
		ValueSize:  expectedSize,
		MaxEntries: 1,
		Name:       "test_v4_state",
	})
}
<commit_msg>Update naming and comment for golang State struct to match C code<commit_after>// Copyright (c) 2020-2021 Tigera, Inc. All rights reserved.

// 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 state

import (
	"unsafe"

	log "github.com/sirupsen/logrus"

	"github.com/projectcalico/felix/bpf"
)

type PolicyResult int32

const (
	PolicyNoMatch PolicyResult = iota
	PolicyAllow
	PolicyDeny
	PolicyTailCallFailed = 10
)

// struct cali_tc_state {
//    __be32 ip_src;
//    __be32 ip_dst;
//    __be32 pre_nat_ip_dst;
//    __be32 post_nat_ip_dst;
//    __be32 tun_ip;
//    __s32 pol_rc;
//    __u16 sport;
//    __u16 dport;
//    __u16 pre_nat_dport;
//    __u16 post_nat_dport;
//    __u8 ip_proto;
//    __u8 flags;
//    __be16 ip_size;
//    struct calico_ct_result ct_result;
//    struct calico_nat_dest nat_dest;
//    __u64 prog_start_time;
// };
type State struct {
	SrcAddr             uint32
	DstAddr             uint32
	PreNATDstAddr       uint32
	PostNATDstAddr      uint32
	TunIP               uint32
	PolicyRC            PolicyResult
	SrcPort             uint16
	DstPort             uint16
	PreNATDstPort       uint16
	PostNATDstPort      uint16
	IPProto             uint8
	Flags               uint8
	IPSize              uint16
	ConntrackRCFlags    uint32
	ConntrackNATIPPort  uint64
	ConntrackTunIP      uint32
	ConntrackIfIndexFwd uint32
	ConntrackIfIndexCtd uint32
	NATData             uint64
	ProgStartTime       uint64
}

const expectedSize = 80

func (s *State) AsBytes() []byte {
	size := unsafe.Sizeof(State{})
	if size != expectedSize {
		log.WithField("size", size).Panic("Incorrect struct size")
	}
	bPtr := (*[expectedSize]byte)(unsafe.Pointer(s))
	bytes := make([]byte, expectedSize)
	copy(bytes, bPtr[:])
	return bytes
}

func StateFromBytes(bytes []byte) State {
	s := State{}
	bPtr := (*[expectedSize]byte)(unsafe.Pointer(&s))
	copy(bPtr[:], bytes)
	return s
}

func Map(mc *bpf.MapContext) bpf.Map {
	return mc.NewPinnedMap(bpf.MapParameters{
		Filename:   "/sys/fs/bpf/tc/globals/cali_v4_state",
		Type:       "percpu_array",
		KeySize:    4,
		ValueSize:  expectedSize,
		MaxEntries: 1,
		Name:       "cali_v4_state",
		Version:    3,
	})
}

func MapForTest(mc *bpf.MapContext) bpf.Map {
	return mc.NewPinnedMap(bpf.MapParameters{
		Filename:   "/sys/fs/bpf/tc/globals/test_v4_state",
		Type:       "array",
		KeySize:    4,
		ValueSize:  expectedSize,
		MaxEntries: 1,
		Name:       "test_v4_state",
	})
}
<|endoftext|>


<commit_before>package raw

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

func TestByteSliceWithNil(t *testing.T) {
	b := ByteSlice(nil)
	buf := ByteSlice(b)
	if len(buf) != 0 {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), 0)
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), 0)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithByteSlice(t *testing.T) {
	b := make([]byte, 10, 10)
	buf := ByteSlice(b)
	if len(buf) != len(b) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(b))
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), cap(b))
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithMap(t *testing.T) {
	m := make(map[int]int)
	defer func() {
		if x := recover(); x == nil {
			t.Fatalf("should have raised a panic")
		}
	}()
	ByteSlice(m)
}

func TestByteSliceWithChannel(t *testing.T) {
	c := make(chan int)
	defer func() {
		if x := recover(); x == nil {
			t.Fatalf("should have raised a panic")
		}
	}()
	ByteSlice(c)
}

func TestByteSliceWithInterface(t *testing.T) {
	t.Log("Awaiting bug fix for incorrect reporting of interface{} value size with unsafe.Sizeof()")
/*	var i interface{} = make([]byte, INTERFACE.size, INTERFACE.size)
	b := i.([]byte)
	buf := ByteSlice(i)
	if len(buf) != len(b) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(b))
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), cap(b))
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
*/}

func TestByteSliceWithString(t *testing.T) {
	s := "hello"
	buf := ByteSlice(s)
	if len(buf) != len(s) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(s))
	}
	if cap(buf) != len(s) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), len(s))
	}

	stringheader := *(*reflect.StringHeader)(unsafe.Pointer(&s))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if stringheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", stringheader.Data, bufheader.Data)
	}
}

type Point struct {
	x			int32
	y			int32
	z			int32
}

func TestByteSliceWithStructValue(t *testing.T) {
	point := Point{ 3, 4, 5 }
	buf := ByteSlice(&point)

	size := int(unsafe.Sizeof(point))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(&point))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

func TestByteSliceWithStructPointer(t *testing.T) {
	point := &Point{ 3, 4, 5 }
	buf := ByteSlice(point)

	size := int(unsafe.Sizeof(*point))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(point))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

type TaggedPoint struct {
	Point
	tag			string
}

func TestByteSliceWithEmbeddedStructValue(t *testing.T) {
	point := Point{ 3, 4, 5 }
	tag := &TaggedPoint{ point, "this is a tag" }
	buf := ByteSlice(tag)

	size := int(unsafe.Sizeof(*tag))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(tag))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

type TaggedPointReference struct {
	*Point
	tag			string
}

func TestByteSliceWithEmbeddedStructPointer(t *testing.T) {
	point := Point{ 3, 4, 5 }
	tag := &TaggedPointReference{ &point, "this is a tag" }
	buf := ByteSlice(tag)

	size := int(unsafe.Sizeof(*tag))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(tag))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

func TestByteSliceWithInt32Slice(t *testing.T) {
	i := make([]int32, 10, 10)
	buf := ByteSlice(i)
	size := len(i) * INT32.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&i))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithInt64Slice(t *testing.T) {
	i := make([]int64, 10, 10)
	buf := ByteSlice(i)
	size := len(i) * INT64.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&i))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithFloat32Slice(t *testing.T) {
	f := make([]float32, 10, 10)
	buf := ByteSlice(f)
	size := len(f) * FLOAT32.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&f))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithFloat64Slice(t *testing.T) {
	f := make([]float64, 10, 10)
	buf := ByteSlice(f)
	size := len(f) * FLOAT64.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&f))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func ValidateNumericByteSlice(t *testing.T, value interface{}) {
	var size	int
	var addr	uintptr
	var numtype	reflect.Type

	v := reflect.ValueOf(value)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	size = int(v.Type().Size())
	addr = v.UnsafeAddr()
	numtype = v.Type()

	buf := ByteSlice(value)
	if len(buf) != size {
		t.Fatalf("%v: byte buffer lengths differ: %v != %v", numtype, len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("%v: byte buffer capacities differ: %v != %v", numtype, cap(buf), size)
	}

	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if addr != bufheader.Data {
		t.Fatalf("%v: addresses don't match: %v != %v", numtype, addr, bufheader.Data)
	}
}

func TestByteSliceWithNumbers(t *testing.T) {
	var i		int
	var i8		int8
	var i16		int16
	var i32 	int32
	var i64		int64
	var u		uint
	var u8		uint8
	var u16		uint16
	var u32 	uint32
	var u64		uint64
	var f32 	float32
	var f64		float64
	var c64		complex64
	var c128	complex64

	ValidateNumericByteSlice(t, &i)
	ValidateNumericByteSlice(t, &i8)
	ValidateNumericByteSlice(t, &i16)
	ValidateNumericByteSlice(t, &i32)
	ValidateNumericByteSlice(t, &i64)
	ValidateNumericByteSlice(t, &u)
	ValidateNumericByteSlice(t, &u8)
	ValidateNumericByteSlice(t, &u16)
	ValidateNumericByteSlice(t, &u32)
	ValidateNumericByteSlice(t, &u64)
	ValidateNumericByteSlice(t, &f32)
	ValidateNumericByteSlice(t, &f64)
	ValidateNumericByteSlice(t, &c64)
	ValidateNumericByteSlice(t, &c128)
}

func TestByteSliceWithNumbersInSlice(t *testing.T) {
	values := []interface{}{	int(0), int8(0), int16(0), int32(0), int64(0),
								uint(0), uint8(0), uint16(0), uint32(0), uint64(0),
								float32(0.0), float64(0.0),
								complex64(0), complex128(0) }
	for i, _ := range values {
		ValidateNumericByteSlice(t, &values[i])
	}
}

func TestByteSliceDataAddress(t *testing.T) {
	DataAddress([]byte{})
	DataAddress(make([]byte, 0, 0))
}

func TestSliceByteCopy(t *testing.T) {
	s := []int{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 }
	for i := 0; i < len(s); i++ {
		c := cap(s) / 2
		if c < i {
			c = i
		}
		d := make([]int, i, c)
		ByteCopy(d, s)
		if len(d) != i {
			t.Fatalf("destination buffer length changed: %v != %v", len(d), i)
		}
		if !reflect.DeepEqual(d, s[:i]) {
			t.Fatalf("buffer contents differ: %v != %v", d, s[:i])
		}
	}
}<commit_msg>added a test for byteslicing structs with zero fields<commit_after>package raw

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

func TestByteSliceWithNil(t *testing.T) {
	b := ByteSlice(nil)
	buf := ByteSlice(b)
	if len(buf) != 0 {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), 0)
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), 0)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithByteSlice(t *testing.T) {
	b := make([]byte, 10, 10)
	buf := ByteSlice(b)
	if len(buf) != len(b) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(b))
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), cap(b))
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithMap(t *testing.T) {
	m := make(map[int]int)
	defer func() {
		if x := recover(); x == nil {
			t.Fatalf("should have raised a panic")
		}
	}()
	ByteSlice(m)
}

func TestByteSliceWithChannel(t *testing.T) {
	c := make(chan int)
	defer func() {
		if x := recover(); x == nil {
			t.Fatalf("should have raised a panic")
		}
	}()
	ByteSlice(c)
}

func TestByteSliceWithInterface(t *testing.T) {
	t.Log("Awaiting bug fix for incorrect reporting of interface{} value size with unsafe.Sizeof()")
/*	var i interface{} = make([]byte, INTERFACE.size, INTERFACE.size)
	b := i.([]byte)
	buf := ByteSlice(i)
	if len(buf) != len(b) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(b))
	}
	if cap(buf) != cap(b) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), cap(b))
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&b))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
*/}

func TestByteSliceWithString(t *testing.T) {
	s := "hello"
	buf := ByteSlice(s)
	if len(buf) != len(s) {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), len(s))
	}
	if cap(buf) != len(s) {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), len(s))
	}

	stringheader := *(*reflect.StringHeader)(unsafe.Pointer(&s))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if stringheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", stringheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithEmptyStructValue(t *testing.T) {
	s := struct {}{}
	buf := ByteSlice(&s)

	size := int(unsafe.Sizeof(struct {}{}))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(&s))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

type Point struct {
	x			int32
	y			int32
	z			int32
}

func TestByteSliceWithStructValue(t *testing.T) {
	point := Point{ 3, 4, 5 }
	buf := ByteSlice(&point)

	size := int(unsafe.Sizeof(point))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(&point))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

func TestByteSliceWithStructPointer(t *testing.T) {
	point := &Point{ 3, 4, 5 }
	buf := ByteSlice(point)

	size := int(unsafe.Sizeof(*point))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(point))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

type TaggedPoint struct {
	Point
	tag			string
}

func TestByteSliceWithEmbeddedStructValue(t *testing.T) {
	point := Point{ 3, 4, 5 }
	tag := &TaggedPoint{ point, "this is a tag" }
	buf := ByteSlice(tag)

	size := int(unsafe.Sizeof(*tag))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(tag))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

type TaggedPointReference struct {
	*Point
	tag			string
}

func TestByteSliceWithEmbeddedStructPointer(t *testing.T) {
	point := Point{ 3, 4, 5 }
	tag := &TaggedPointReference{ &point, "this is a tag" }
	buf := ByteSlice(tag)

	size := int(unsafe.Sizeof(*tag))
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	base_address := uintptr(unsafe.Pointer(tag))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if base_address != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", base_address, bufheader.Data)
	}
}

func TestByteSliceWithInt32Slice(t *testing.T) {
	i := make([]int32, 10, 10)
	buf := ByteSlice(i)
	size := len(i) * INT32.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&i))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithInt64Slice(t *testing.T) {
	i := make([]int64, 10, 10)
	buf := ByteSlice(i)
	size := len(i) * INT64.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&i))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithFloat32Slice(t *testing.T) {
	f := make([]float32, 10, 10)
	buf := ByteSlice(f)
	size := len(f) * FLOAT32.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&f))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func TestByteSliceWithFloat64Slice(t *testing.T) {
	f := make([]float64, 10, 10)
	buf := ByteSlice(f)
	size := len(f) * FLOAT64.size
	if len(buf) != size {
		t.Fatalf("byte buffer lengths differ: %v != %v", len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("byte buffer capacities differ: %v != %v", cap(buf), size)
	}

	sliceheader := *(*reflect.SliceHeader)(unsafe.Pointer(&f))
	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if sliceheader.Data != bufheader.Data {
		t.Fatalf("slice addresses don't match: %v != %v", sliceheader.Data, bufheader.Data)
	}
}

func ValidateNumericByteSlice(t *testing.T, value interface{}) {
	var size	int
	var addr	uintptr
	var numtype	reflect.Type

	v := reflect.ValueOf(value)
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}
	size = int(v.Type().Size())
	addr = v.UnsafeAddr()
	numtype = v.Type()

	buf := ByteSlice(value)
	if len(buf) != size {
		t.Fatalf("%v: byte buffer lengths differ: %v != %v", numtype, len(buf), size)
	}

	if cap(buf) != size {
		t.Fatalf("%v: byte buffer capacities differ: %v != %v", numtype, cap(buf), size)
	}

	bufheader := *(*reflect.SliceHeader)(unsafe.Pointer(&buf))
	if addr != bufheader.Data {
		t.Fatalf("%v: addresses don't match: %v != %v", numtype, addr, bufheader.Data)
	}
}

func TestByteSliceWithNumbers(t *testing.T) {
	var i		int
	var i8		int8
	var i16		int16
	var i32 	int32
	var i64		int64
	var u		uint
	var u8		uint8
	var u16		uint16
	var u32 	uint32
	var u64		uint64
	var f32 	float32
	var f64		float64
	var c64		complex64
	var c128	complex64

	ValidateNumericByteSlice(t, &i)
	ValidateNumericByteSlice(t, &i8)
	ValidateNumericByteSlice(t, &i16)
	ValidateNumericByteSlice(t, &i32)
	ValidateNumericByteSlice(t, &i64)
	ValidateNumericByteSlice(t, &u)
	ValidateNumericByteSlice(t, &u8)
	ValidateNumericByteSlice(t, &u16)
	ValidateNumericByteSlice(t, &u32)
	ValidateNumericByteSlice(t, &u64)
	ValidateNumericByteSlice(t, &f32)
	ValidateNumericByteSlice(t, &f64)
	ValidateNumericByteSlice(t, &c64)
	ValidateNumericByteSlice(t, &c128)
}

func TestByteSliceWithNumbersInSlice(t *testing.T) {
	values := []interface{}{	int(0), int8(0), int16(0), int32(0), int64(0),
								uint(0), uint8(0), uint16(0), uint32(0), uint64(0),
								float32(0.0), float64(0.0),
								complex64(0), complex128(0) }
	for i, _ := range values {
		ValidateNumericByteSlice(t, &values[i])
	}
}

func TestByteSliceDataAddress(t *testing.T) {
	DataAddress([]byte{})
	DataAddress(make([]byte, 0, 0))
}

func TestSliceByteCopy(t *testing.T) {
	s := []int{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 }
	for i := 0; i < len(s); i++ {
		c := cap(s) / 2
		if c < i {
			c = i
		}
		d := make([]int, i, c)
		ByteCopy(d, s)
		if len(d) != i {
			t.Fatalf("destination buffer length changed: %v != %v", len(d), i)
		}
		if !reflect.DeepEqual(d, s[:i]) {
			t.Fatalf("buffer contents differ: %v != %v", d, s[:i])
		}
	}
}<|endoftext|>


<commit_before>package basichost

import (
	"bytes"
	"context"
	"io"
	"sort"
	"testing"
	"time"

	testutil "github.com/libp2p/go-testutil"

	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	protocol "github.com/libp2p/go-libp2p-protocol"
	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
	ma "github.com/multiformats/go-multiaddr"
	madns "github.com/multiformats/go-multiaddr-dns"
)

func TestHostSimple(t *testing.T) {

	ctx := context.Background()
	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))
	defer h1.Close()
	defer h2.Close()

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	piper, pipew := io.Pipe()
	h2.SetStreamHandler(protocol.TestingID, func(s inet.Stream) {
		defer s.Close()
		w := io.MultiWriter(s, pipew)
		io.Copy(w, s) // mirror everything
	})

	s, err := h1.NewStream(ctx, h2pi.ID, protocol.TestingID)
	if err != nil {
		t.Fatal(err)
	}

	// write to the stream
	buf1 := []byte("abcdefghijkl")
	if _, err := s.Write(buf1); err != nil {
		t.Fatal(err)
	}

	// get it from the stream (echoed)
	buf2 := make([]byte, len(buf1))
	if _, err := io.ReadFull(s, buf2); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(buf1, buf2) {
		t.Fatalf("buf1 != buf2 -- %x != %x", buf1, buf2)
	}

	// get it from the pipe (tee)
	buf3 := make([]byte, len(buf1))
	if _, err := io.ReadFull(piper, buf3); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(buf1, buf3) {
		t.Fatalf("buf1 != buf3 -- %x != %x", buf1, buf3)
	}
}

func TestHostAddrsFactory(t *testing.T) {
	maddr := ma.StringCast("/ip4/1.2.3.4/tcp/1234")
	addrsFactory := func(addrs []ma.Multiaddr) []ma.Multiaddr {
		return []ma.Multiaddr{maddr}
	}

	ctx := context.Background()
	h := New(swarmt.GenSwarm(t, ctx), AddrsFactory(addrsFactory))
	defer h.Close()

	addrs := h.Addrs()
	if len(addrs) != 1 {
		t.Fatalf("expected 1 addr, got %d", len(addrs))
	}
	if addrs[0] != maddr {
		t.Fatalf("expected %s, got %s", maddr.String(), addrs[0].String())
	}
}

func getHostPair(ctx context.Context, t *testing.T) (host.Host, host.Host) {
	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	return h1, h2
}

func assertWait(t *testing.T, c chan protocol.ID, exp protocol.ID) {
	select {
	case proto := <-c:
		if proto != exp {
			t.Fatal("should have connected on ", exp)
		}
	case <-time.After(time.Second * 5):
		t.Fatal("timeout waiting for stream")
	}
}

func TestHostProtoPreference(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	protoOld := protocol.ID("/testing")
	protoNew := protocol.ID("/testing/1.1.0")
	protoMinor := protocol.ID("/testing/1.2.0")

	connectedOn := make(chan protocol.ID)

	handler := func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	}

	h1.SetStreamHandler(protoOld, handler)

	s, err := h2.NewStream(ctx, h1.ID(), protoMinor, protoNew, protoOld)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoOld)
	s.Close()

	mfunc, err := host.MultistreamSemverMatcher(protoMinor)
	if err != nil {
		t.Fatal(err)
	}

	h1.SetStreamHandlerMatch(protoMinor, mfunc, handler)

	// remembered preference will be chosen first, even when the other side newly supports it
	s2, err := h2.NewStream(ctx, h1.ID(), protoMinor, protoNew, protoOld)
	if err != nil {
		t.Fatal(err)
	}

	// required to force 'lazy' handshake
	_, err = s2.Write([]byte("hello"))
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoOld)

	s2.Close()

	s3, err := h2.NewStream(ctx, h1.ID(), protoMinor)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoMinor)
	s3.Close()
}

func TestHostProtoMismatch(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	h1.SetStreamHandler("/super", func(s inet.Stream) {
		t.Error("shouldnt get here")
		s.Reset()
	})

	_, err := h2.NewStream(ctx, h1.ID(), "/foo", "/bar", "/baz/1.0.0")
	if err == nil {
		t.Fatal("expected new stream to fail")
	}
}

func TestHostProtoPreknowledge(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))

	conn := make(chan protocol.ID)
	handler := func(s inet.Stream) {
		conn <- s.Protocol()
		s.Close()
	}

	h1.SetStreamHandler("/super", handler)

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}
	defer h1.Close()
	defer h2.Close()

	// wait for identify handshake to finish completely
	select {
	case <-h1.ids.IdentifyWait(h1.Network().ConnsToPeer(h2.ID())[0]):
	case <-time.After(time.Second * 5):
		t.Fatal("timed out waiting for identify")
	}

	select {
	case <-h2.ids.IdentifyWait(h2.Network().ConnsToPeer(h1.ID())[0]):
	case <-time.After(time.Second * 5):
		t.Fatal("timed out waiting for identify")
	}

	h1.SetStreamHandler("/foo", handler)

	s, err := h2.NewStream(ctx, h1.ID(), "/foo", "/bar", "/super")
	if err != nil {
		t.Fatal(err)
	}

	select {
	case p := <-conn:
		t.Fatal("shouldnt have gotten connection yet, we should have a lazy stream: ", p)
	case <-time.After(time.Millisecond * 50):
	}

	_, err = s.Read(nil)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, conn, "/super")

	s.Close()
}

func TestNewDialOld(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	connectedOn := make(chan protocol.ID)
	h1.SetStreamHandler("/testing", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	s, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing")

	if s.Protocol() != "/testing" {
		t.Fatal("shoould have gotten /testing")
	}

	s.Close()
}

func TestProtoDowngrade(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	connectedOn := make(chan protocol.ID)
	h1.SetStreamHandler("/testing/1.0.0", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	s, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing/1.0.0")

	if s.Protocol() != "/testing/1.0.0" {
		t.Fatal("shoould have gotten /testing")
	}
	s.Close()

	h1.Network().ConnsToPeer(h2.ID())[0].Close()

	time.Sleep(time.Millisecond * 50) // allow notifications to propogate
	h1.RemoveStreamHandler("/testing/1.0.0")
	h1.SetStreamHandler("/testing", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	s2, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	_, err = s2.Write(nil)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing")

	if s2.Protocol() != "/testing" {
		t.Fatal("shoould have gotten /testing")
	}
	s2.Close()

}

func TestAddrResolution(t *testing.T) {
	ctx := context.Background()

	p1, err := testutil.RandPeerID()
	if err != nil {
		t.Error(err)
	}
	p2, err := testutil.RandPeerID()
	if err != nil {
		t.Error(err)
	}
	addr1 := ma.StringCast("/dnsaddr/example.com")
	addr2 := ma.StringCast("/ip4/192.0.2.1/tcp/123")
	p2paddr1 := ma.StringCast("/dnsaddr/example.com/ipfs/" + p1.Pretty())
	p2paddr2 := ma.StringCast("/ip4/192.0.2.1/tcp/123/ipfs/" + p1.Pretty())
	p2paddr3 := ma.StringCast("/ip4/192.0.2.1/tcp/123/ipfs/" + p2.Pretty())

	backend := &madns.MockBackend{
		TXT: map[string][]string{"_dnsaddr.example.com": []string{
			"dnsaddr=" + p2paddr2.String(), "dnsaddr=" + p2paddr3.String(),
		}},
	}
	resolver := &madns.Resolver{Backend: backend}

	h := New(swarmt.GenSwarm(t, ctx), resolver)
	defer h.Close()

	pi, err := pstore.InfoFromP2pAddr(p2paddr1)
	if err != nil {
		t.Error(err)
	}

	tctx, cancel := context.WithTimeout(ctx, time.Millisecond*100)
	defer cancel()
	_ = h.Connect(tctx, *pi)

	addrs := h.Peerstore().Addrs(pi.ID)
	sort.Sort(sortedMultiaddrs(addrs))

	if len(addrs) != 2 || !addrs[0].Equal(addr1) || !addrs[1].Equal(addr2) {
		t.Fatalf("expected [%s %s], got %+v", addr1, addr2, addrs)
	}
}

type sortedMultiaddrs []ma.Multiaddr

func (sma sortedMultiaddrs) Len() int      { return len(sma) }
func (sma sortedMultiaddrs) Swap(i, j int) { sma[i], sma[j] = sma[j], sma[i] }
func (sma sortedMultiaddrs) Less(i, j int) bool {
	return bytes.Compare(sma[i].Bytes(), sma[j].Bytes()) == 1
}
<commit_msg>fix buggy addr equal in test<commit_after>package basichost

import (
	"bytes"
	"context"
	"io"
	"sort"
	"testing"
	"time"

	testutil "github.com/libp2p/go-testutil"

	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	protocol "github.com/libp2p/go-libp2p-protocol"
	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
	ma "github.com/multiformats/go-multiaddr"
	madns "github.com/multiformats/go-multiaddr-dns"
)

func TestHostSimple(t *testing.T) {

	ctx := context.Background()
	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))
	defer h1.Close()
	defer h2.Close()

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	piper, pipew := io.Pipe()
	h2.SetStreamHandler(protocol.TestingID, func(s inet.Stream) {
		defer s.Close()
		w := io.MultiWriter(s, pipew)
		io.Copy(w, s) // mirror everything
	})

	s, err := h1.NewStream(ctx, h2pi.ID, protocol.TestingID)
	if err != nil {
		t.Fatal(err)
	}

	// write to the stream
	buf1 := []byte("abcdefghijkl")
	if _, err := s.Write(buf1); err != nil {
		t.Fatal(err)
	}

	// get it from the stream (echoed)
	buf2 := make([]byte, len(buf1))
	if _, err := io.ReadFull(s, buf2); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(buf1, buf2) {
		t.Fatalf("buf1 != buf2 -- %x != %x", buf1, buf2)
	}

	// get it from the pipe (tee)
	buf3 := make([]byte, len(buf1))
	if _, err := io.ReadFull(piper, buf3); err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(buf1, buf3) {
		t.Fatalf("buf1 != buf3 -- %x != %x", buf1, buf3)
	}
}

func TestHostAddrsFactory(t *testing.T) {
	maddr := ma.StringCast("/ip4/1.2.3.4/tcp/1234")
	addrsFactory := func(addrs []ma.Multiaddr) []ma.Multiaddr {
		return []ma.Multiaddr{maddr}
	}

	ctx := context.Background()
	h := New(swarmt.GenSwarm(t, ctx), AddrsFactory(addrsFactory))
	defer h.Close()

	addrs := h.Addrs()
	if len(addrs) != 1 {
		t.Fatalf("expected 1 addr, got %d", len(addrs))
	}
	if !addrs[0].Equal(maddr) {
		t.Fatalf("expected %s, got %s", maddr.String(), addrs[0].String())
	}
}

func getHostPair(ctx context.Context, t *testing.T) (host.Host, host.Host) {
	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	return h1, h2
}

func assertWait(t *testing.T, c chan protocol.ID, exp protocol.ID) {
	select {
	case proto := <-c:
		if proto != exp {
			t.Fatal("should have connected on ", exp)
		}
	case <-time.After(time.Second * 5):
		t.Fatal("timeout waiting for stream")
	}
}

func TestHostProtoPreference(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	protoOld := protocol.ID("/testing")
	protoNew := protocol.ID("/testing/1.1.0")
	protoMinor := protocol.ID("/testing/1.2.0")

	connectedOn := make(chan protocol.ID)

	handler := func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	}

	h1.SetStreamHandler(protoOld, handler)

	s, err := h2.NewStream(ctx, h1.ID(), protoMinor, protoNew, protoOld)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoOld)
	s.Close()

	mfunc, err := host.MultistreamSemverMatcher(protoMinor)
	if err != nil {
		t.Fatal(err)
	}

	h1.SetStreamHandlerMatch(protoMinor, mfunc, handler)

	// remembered preference will be chosen first, even when the other side newly supports it
	s2, err := h2.NewStream(ctx, h1.ID(), protoMinor, protoNew, protoOld)
	if err != nil {
		t.Fatal(err)
	}

	// required to force 'lazy' handshake
	_, err = s2.Write([]byte("hello"))
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoOld)

	s2.Close()

	s3, err := h2.NewStream(ctx, h1.ID(), protoMinor)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, protoMinor)
	s3.Close()
}

func TestHostProtoMismatch(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	h1.SetStreamHandler("/super", func(s inet.Stream) {
		t.Error("shouldnt get here")
		s.Reset()
	})

	_, err := h2.NewStream(ctx, h1.ID(), "/foo", "/bar", "/baz/1.0.0")
	if err == nil {
		t.Fatal("expected new stream to fail")
	}
}

func TestHostProtoPreknowledge(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1 := New(swarmt.GenSwarm(t, ctx))
	h2 := New(swarmt.GenSwarm(t, ctx))

	conn := make(chan protocol.ID)
	handler := func(s inet.Stream) {
		conn <- s.Protocol()
		s.Close()
	}

	h1.SetStreamHandler("/super", handler)

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}
	defer h1.Close()
	defer h2.Close()

	// wait for identify handshake to finish completely
	select {
	case <-h1.ids.IdentifyWait(h1.Network().ConnsToPeer(h2.ID())[0]):
	case <-time.After(time.Second * 5):
		t.Fatal("timed out waiting for identify")
	}

	select {
	case <-h2.ids.IdentifyWait(h2.Network().ConnsToPeer(h1.ID())[0]):
	case <-time.After(time.Second * 5):
		t.Fatal("timed out waiting for identify")
	}

	h1.SetStreamHandler("/foo", handler)

	s, err := h2.NewStream(ctx, h1.ID(), "/foo", "/bar", "/super")
	if err != nil {
		t.Fatal(err)
	}

	select {
	case p := <-conn:
		t.Fatal("shouldnt have gotten connection yet, we should have a lazy stream: ", p)
	case <-time.After(time.Millisecond * 50):
	}

	_, err = s.Read(nil)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, conn, "/super")

	s.Close()
}

func TestNewDialOld(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	connectedOn := make(chan protocol.ID)
	h1.SetStreamHandler("/testing", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	s, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing")

	if s.Protocol() != "/testing" {
		t.Fatal("shoould have gotten /testing")
	}

	s.Close()
}

func TestProtoDowngrade(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	h1, h2 := getHostPair(ctx, t)
	defer h1.Close()
	defer h2.Close()

	connectedOn := make(chan protocol.ID)
	h1.SetStreamHandler("/testing/1.0.0", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	s, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing/1.0.0")

	if s.Protocol() != "/testing/1.0.0" {
		t.Fatal("shoould have gotten /testing")
	}
	s.Close()

	h1.Network().ConnsToPeer(h2.ID())[0].Close()

	time.Sleep(time.Millisecond * 50) // allow notifications to propogate
	h1.RemoveStreamHandler("/testing/1.0.0")
	h1.SetStreamHandler("/testing", func(s inet.Stream) {
		connectedOn <- s.Protocol()
		s.Close()
	})

	h2pi := h2.Peerstore().PeerInfo(h2.ID())
	if err := h1.Connect(ctx, h2pi); err != nil {
		t.Fatal(err)
	}

	s2, err := h2.NewStream(ctx, h1.ID(), "/testing/1.0.0", "/testing")
	if err != nil {
		t.Fatal(err)
	}

	_, err = s2.Write(nil)
	if err != nil {
		t.Fatal(err)
	}

	assertWait(t, connectedOn, "/testing")

	if s2.Protocol() != "/testing" {
		t.Fatal("shoould have gotten /testing")
	}
	s2.Close()

}

func TestAddrResolution(t *testing.T) {
	ctx := context.Background()

	p1, err := testutil.RandPeerID()
	if err != nil {
		t.Error(err)
	}
	p2, err := testutil.RandPeerID()
	if err != nil {
		t.Error(err)
	}
	addr1 := ma.StringCast("/dnsaddr/example.com")
	addr2 := ma.StringCast("/ip4/192.0.2.1/tcp/123")
	p2paddr1 := ma.StringCast("/dnsaddr/example.com/ipfs/" + p1.Pretty())
	p2paddr2 := ma.StringCast("/ip4/192.0.2.1/tcp/123/ipfs/" + p1.Pretty())
	p2paddr3 := ma.StringCast("/ip4/192.0.2.1/tcp/123/ipfs/" + p2.Pretty())

	backend := &madns.MockBackend{
		TXT: map[string][]string{"_dnsaddr.example.com": []string{
			"dnsaddr=" + p2paddr2.String(), "dnsaddr=" + p2paddr3.String(),
		}},
	}
	resolver := &madns.Resolver{Backend: backend}

	h := New(swarmt.GenSwarm(t, ctx), resolver)
	defer h.Close()

	pi, err := pstore.InfoFromP2pAddr(p2paddr1)
	if err != nil {
		t.Error(err)
	}

	tctx, cancel := context.WithTimeout(ctx, time.Millisecond*100)
	defer cancel()
	_ = h.Connect(tctx, *pi)

	addrs := h.Peerstore().Addrs(pi.ID)
	sort.Sort(sortedMultiaddrs(addrs))

	if len(addrs) != 2 || !addrs[0].Equal(addr1) || !addrs[1].Equal(addr2) {
		t.Fatalf("expected [%s %s], got %+v", addr1, addr2, addrs)
	}
}

type sortedMultiaddrs []ma.Multiaddr

func (sma sortedMultiaddrs) Len() int      { return len(sma) }
func (sma sortedMultiaddrs) Swap(i, j int) { sma[i], sma[j] = sma[j], sma[i] }
func (sma sortedMultiaddrs) Less(i, j int) bool {
	return bytes.Compare(sma[i].Bytes(), sma[j].Bytes()) == 1
}
<|endoftext|>


<commit_before>package stream_test

import (
	"context"
	"errors"
	"net"
	"sync"
	"time"

	insecure "github.com/libp2p/go-conn-security/insecure"
	peer "github.com/libp2p/go-libp2p-peer"
	tpt "github.com/libp2p/go-libp2p-transport"
	st "github.com/libp2p/go-libp2p-transport-upgrader"
	smux "github.com/libp2p/go-stream-muxer"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	yamux "github.com/whyrusleeping/go-smux-yamux"

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

// negotiatingMuxer sets up a new yamux connection
// It makes sure that this happens at the same time for client and server.
type negotiatingMuxer struct{}

func (m *negotiatingMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	var err error
	// run a fake muxer negotiation
	if isServer {
		_, err = c.Write([]byte("setup"))
	} else {
		_, err = c.Read(make([]byte, 5))
	}
	if err != nil {
		return nil, err
	}
	return yamux.DefaultTransport.NewConn(c, isServer)
}

// blockingMuxer blocks the muxer negotiation until the contain chan is closed
type blockingMuxer struct {
	unblock chan struct{}
}

var _ smux.Transport = &blockingMuxer{}

func newBlockingMuxer() *blockingMuxer { return &blockingMuxer{unblock: make(chan struct{})} }
func (m *blockingMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	<-m.unblock
	return (&negotiatingMuxer{}).NewConn(c, isServer)
}
func (m *blockingMuxer) Unblock() { close(m.unblock) }

// errorMuxer is a muxer that errors while setting up
type errorMuxer struct{}

var _ smux.Transport = &errorMuxer{}

func (m *errorMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	return nil, errors.New("mux error")
}

var _ = Describe("Listener", func() {
	var (
		defaultUpgrader = &st.Upgrader{
			Secure: insecure.New(peer.ID(1)),
			Muxer:  &negotiatingMuxer{},
		}
	)

	testConn := func(clientConn, serverConn tpt.Conn) {
		cstr, err := clientConn.OpenStream()
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		_, err = cstr.Write([]byte("foobar"))
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		sstr, err := serverConn.AcceptStream()
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		b := make([]byte, 6)
		_, err = sstr.Read(b)
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		ExpectWithOffset(0, b).To(Equal([]byte("foobar")))
	}

	createListener := func(upgrader *st.Upgrader) tpt.Listener {
		addr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		ln, err := manet.Listen(addr)
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		return upgrader.UpgradeListener(nil, ln)
	}

	dial := func(upgrader *st.Upgrader, raddr ma.Multiaddr, p peer.ID) (tpt.Conn, error) {
		macon, err := manet.Dial(raddr)
		if err != nil {
			return nil, err
		}
		return upgrader.UpgradeOutbound(context.Background(), nil, macon, p)
	}

	BeforeEach(func() {
		tpt.AcceptTimeout = time.Hour
	})

	It("accepts a single connection", func() {
		ln := createListener(defaultUpgrader)
		cconn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
		Expect(err).ToNot(HaveOccurred())
		sconn, err := ln.Accept()
		Expect(err).ToNot(HaveOccurred())
		testConn(cconn, sconn)
	})

	It("accepts multiple connections", func() {
		ln := createListener(defaultUpgrader)
		const num = 10
		for i := 0; i < 10; i++ {
			cconn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
			Expect(err).ToNot(HaveOccurred())
			sconn, err := ln.Accept()
			Expect(err).ToNot(HaveOccurred())
			testConn(cconn, sconn)
		}
	})

	It("closes connections if they are not accepted", func() {
		const timeout = 200 * time.Millisecond
		tpt.AcceptTimeout = timeout
		ln := createListener(defaultUpgrader)
		conn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
		Expect(err).ToNot(HaveOccurred())
		done := make(chan struct{})
		go func() {
			defer GinkgoRecover()
			str, err := conn.OpenStream()
			Expect(err).ToNot(HaveOccurred())
			// start a Read. It will block until the connection is closed
			str.Read([]byte{0})
			close(done)
		}()
		Consistently(done, timeout/2).ShouldNot(BeClosed())
		Eventually(done, timeout).Should(BeClosed())
	})

	It("doesn't accept connections that fail to setup", func() {
		upgrader := &st.Upgrader{
			Secure: insecure.New(peer.ID(1)),
			Muxer:  &errorMuxer{},
		}
		ln := createListener(upgrader)
		done := make(chan struct{})
		go func() {
			defer GinkgoRecover()
			_, _ = ln.Accept()
			close(done)
		}()
		_, _ = dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
		Consistently(done).ShouldNot(BeClosed())
		// make the goroutine return
		ln.Close()
		Eventually(done).Should(BeClosed())
	})

	Context("concurrency", func() {
		It("sets up connections concurrently", func() {
			num := 3 * st.AcceptQueueLength
			bm := newBlockingMuxer()
			upgrader := &st.Upgrader{
				Secure: insecure.New(peer.ID(1)),
				Muxer:  bm,
			}
			ln := createListener(upgrader)
			accepted := make(chan tpt.Conn, num)
			go func() {
				defer GinkgoRecover()
				for {
					conn, err := ln.Accept()
					if err != nil {
						return
					}
					accepted <- conn
				}
			}()
			var wg sync.WaitGroup
			// start num dials, which all block while setting up the muxer
			for i := 0; i < num; i++ {
				wg.Add(1)
				go func() {
					defer GinkgoRecover()
					_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
					Expect(err).ToNot(HaveOccurred())
					wg.Done()
				}()
			}
			// the dials are still blocked, so we shouldn't have any connection available yet
			Consistently(accepted).Should(BeEmpty())
			bm.Unblock() // make all dials succeed
			Eventually(accepted).Should(HaveLen(num))
			wg.Wait()
		})

		It("stops setting up when the more than AcceptQueueLength connections are waiting to get accepted", func() {
			ln := createListener(defaultUpgrader)
			// setup AcceptQueueLength connections, but don't accept any of them
			dialed := make(chan struct{}, 10*st.AcceptQueueLength) // used as a thread-safe counter
			for i := 0; i < st.AcceptQueueLength; i++ {
				go func() {
					defer GinkgoRecover()
					_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
					Expect(err).ToNot(HaveOccurred())
					dialed <- struct{}{}
				}()
			}
			Eventually(dialed).Should(HaveLen(st.AcceptQueueLength))
			// dial a new connection. This connection should not complete setup, since the queue is full
			go func() {
				defer GinkgoRecover()
				_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
				Expect(err).ToNot(HaveOccurred())
				dialed <- struct{}{}
			}()
			Consistently(dialed).Should(HaveLen(st.AcceptQueueLength))
			// accept a single connection. Now the new connection should be set up, and fill the queue again
			_, err := ln.Accept()
			Expect(err).ToNot(HaveOccurred())
			Eventually(dialed).Should(HaveLen(st.AcceptQueueLength + 1))
		})
	})

	Context("closing", func() {
		It("unblocks Accept when it is closed", func() {
			ln := createListener(defaultUpgrader)
			done := make(chan struct{})
			go func() {
				defer GinkgoRecover()
				_, err := ln.Accept()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("use of closed network connection"))
				close(done)
			}()
			Consistently(done).ShouldNot(BeClosed())
			Expect(ln.Close()).To(Succeed())
			Eventually(done).Should(BeClosed())
		})

		It("doesn't accept new connections when it is closed", func() {
			ln := createListener(defaultUpgrader)
			Expect(ln.Close()).To(Succeed())
			_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
			Expect(err).To(HaveOccurred())
		})

		It("closes incoming connections that have not yet been accepted", func() {
			ln := createListener(defaultUpgrader)
			conn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
			Expect(conn.IsClosed()).To(BeFalse())
			Expect(err).ToNot(HaveOccurred())
			Expect(ln.Close()).To(Succeed())
			Eventually(conn.IsClosed).Should(BeTrue())
		})
	})
})
<commit_msg>use mplex for tests instead of yamux<commit_after>package stream_test

import (
	"context"
	"errors"
	"net"
	"sync"
	"time"

	insecure "github.com/libp2p/go-conn-security/insecure"
	peer "github.com/libp2p/go-libp2p-peer"
	tpt "github.com/libp2p/go-libp2p-transport"
	st "github.com/libp2p/go-libp2p-transport-upgrader"
	smux "github.com/libp2p/go-stream-muxer"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mplex "github.com/whyrusleeping/go-smux-multiplex"

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

// negotiatingMuxer sets up a new mplex connection
// It makes sure that this happens at the same time for client and server.
type negotiatingMuxer struct{}

func (m *negotiatingMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	var err error
	// run a fake muxer negotiation
	if isServer {
		_, err = c.Write([]byte("setup"))
	} else {
		_, err = c.Read(make([]byte, 5))
	}
	if err != nil {
		return nil, err
	}
	return mplex.DefaultTransport.NewConn(c, isServer)
}

// blockingMuxer blocks the muxer negotiation until the contain chan is closed
type blockingMuxer struct {
	unblock chan struct{}
}

var _ smux.Transport = &blockingMuxer{}

func newBlockingMuxer() *blockingMuxer { return &blockingMuxer{unblock: make(chan struct{})} }
func (m *blockingMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	<-m.unblock
	return (&negotiatingMuxer{}).NewConn(c, isServer)
}
func (m *blockingMuxer) Unblock() { close(m.unblock) }

// errorMuxer is a muxer that errors while setting up
type errorMuxer struct{}

var _ smux.Transport = &errorMuxer{}

func (m *errorMuxer) NewConn(c net.Conn, isServer bool) (smux.Conn, error) {
	return nil, errors.New("mux error")
}

var _ = Describe("Listener", func() {
	var (
		defaultUpgrader = &st.Upgrader{
			Secure: insecure.New(peer.ID(1)),
			Muxer:  &negotiatingMuxer{},
		}
	)

	testConn := func(clientConn, serverConn tpt.Conn) {
		cstr, err := clientConn.OpenStream()
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		_, err = cstr.Write([]byte("foobar"))
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		sstr, err := serverConn.AcceptStream()
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		b := make([]byte, 6)
		_, err = sstr.Read(b)
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		ExpectWithOffset(0, b).To(Equal([]byte("foobar")))
	}

	createListener := func(upgrader *st.Upgrader) tpt.Listener {
		addr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		ln, err := manet.Listen(addr)
		ExpectWithOffset(0, err).ToNot(HaveOccurred())
		return upgrader.UpgradeListener(nil, ln)
	}

	dial := func(upgrader *st.Upgrader, raddr ma.Multiaddr, p peer.ID) (tpt.Conn, error) {
		macon, err := manet.Dial(raddr)
		if err != nil {
			return nil, err
		}
		return upgrader.UpgradeOutbound(context.Background(), nil, macon, p)
	}

	BeforeEach(func() {
		tpt.AcceptTimeout = time.Hour
	})

	It("accepts a single connection", func() {
		ln := createListener(defaultUpgrader)
		cconn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
		Expect(err).ToNot(HaveOccurred())
		sconn, err := ln.Accept()
		Expect(err).ToNot(HaveOccurred())
		testConn(cconn, sconn)
	})

	It("accepts multiple connections", func() {
		ln := createListener(defaultUpgrader)
		const num = 10
		for i := 0; i < 10; i++ {
			cconn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
			Expect(err).ToNot(HaveOccurred())
			sconn, err := ln.Accept()
			Expect(err).ToNot(HaveOccurred())
			testConn(cconn, sconn)
		}
	})

	It("closes connections if they are not accepted", func() {
		const timeout = 200 * time.Millisecond
		tpt.AcceptTimeout = timeout
		ln := createListener(defaultUpgrader)
		conn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
		Expect(err).ToNot(HaveOccurred())
		done := make(chan struct{})
		go func() {
			defer GinkgoRecover()
			str, err := conn.OpenStream()
			Expect(err).ToNot(HaveOccurred())
			// start a Read. It will block until the connection is closed
			str.Read([]byte{0})
			close(done)
		}()
		Consistently(done, timeout/2).ShouldNot(BeClosed())
		Eventually(done, timeout).Should(BeClosed())
	})

	It("doesn't accept connections that fail to setup", func() {
		upgrader := &st.Upgrader{
			Secure: insecure.New(peer.ID(1)),
			Muxer:  &errorMuxer{},
		}
		ln := createListener(upgrader)
		done := make(chan struct{})
		go func() {
			defer GinkgoRecover()
			_, _ = ln.Accept()
			close(done)
		}()
		_, _ = dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
		Consistently(done).ShouldNot(BeClosed())
		// make the goroutine return
		ln.Close()
		Eventually(done).Should(BeClosed())
	})

	Context("concurrency", func() {
		It("sets up connections concurrently", func() {
			num := 3 * st.AcceptQueueLength
			bm := newBlockingMuxer()
			upgrader := &st.Upgrader{
				Secure: insecure.New(peer.ID(1)),
				Muxer:  bm,
			}
			ln := createListener(upgrader)
			accepted := make(chan tpt.Conn, num)
			go func() {
				defer GinkgoRecover()
				for {
					conn, err := ln.Accept()
					if err != nil {
						return
					}
					accepted <- conn
				}
			}()
			var wg sync.WaitGroup
			// start num dials, which all block while setting up the muxer
			for i := 0; i < num; i++ {
				wg.Add(1)
				go func() {
					defer GinkgoRecover()
					_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
					Expect(err).ToNot(HaveOccurred())
					wg.Done()
				}()
			}
			// the dials are still blocked, so we shouldn't have any connection available yet
			Consistently(accepted).Should(BeEmpty())
			bm.Unblock() // make all dials succeed
			Eventually(accepted).Should(HaveLen(num))
			wg.Wait()
		})

		It("stops setting up when the more than AcceptQueueLength connections are waiting to get accepted", func() {
			ln := createListener(defaultUpgrader)
			// setup AcceptQueueLength connections, but don't accept any of them
			dialed := make(chan struct{}, 10*st.AcceptQueueLength) // used as a thread-safe counter
			for i := 0; i < st.AcceptQueueLength; i++ {
				go func() {
					defer GinkgoRecover()
					_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
					Expect(err).ToNot(HaveOccurred())
					dialed <- struct{}{}
				}()
			}
			Eventually(dialed).Should(HaveLen(st.AcceptQueueLength))
			// dial a new connection. This connection should not complete setup, since the queue is full
			go func() {
				defer GinkgoRecover()
				_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
				Expect(err).ToNot(HaveOccurred())
				dialed <- struct{}{}
			}()
			Consistently(dialed).Should(HaveLen(st.AcceptQueueLength))
			// accept a single connection. Now the new connection should be set up, and fill the queue again
			_, err := ln.Accept()
			Expect(err).ToNot(HaveOccurred())
			Eventually(dialed).Should(HaveLen(st.AcceptQueueLength + 1))
		})
	})

	Context("closing", func() {
		It("unblocks Accept when it is closed", func() {
			ln := createListener(defaultUpgrader)
			done := make(chan struct{})
			go func() {
				defer GinkgoRecover()
				_, err := ln.Accept()
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("use of closed network connection"))
				close(done)
			}()
			Consistently(done).ShouldNot(BeClosed())
			Expect(ln.Close()).To(Succeed())
			Eventually(done).Should(BeClosed())
		})

		It("doesn't accept new connections when it is closed", func() {
			ln := createListener(defaultUpgrader)
			Expect(ln.Close()).To(Succeed())
			_, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(1))
			Expect(err).To(HaveOccurred())
		})

		It("closes incoming connections that have not yet been accepted", func() {
			ln := createListener(defaultUpgrader)
			conn, err := dial(defaultUpgrader, ln.Multiaddr(), peer.ID(2))
			Expect(conn.IsClosed()).To(BeFalse())
			Expect(err).ToNot(HaveOccurred())
			Expect(ln.Close()).To(Succeed())
			Eventually(conn.IsClosed).Should(BeTrue())
		})
	})
})
<|endoftext|>


<commit_before>package main // import "honnef.co/go/unused/cmd/unused"

import (
	"flag"
	"fmt"
	"log"
	"os"

	"honnef.co/go/lint/lintutil"
	"honnef.co/go/unused"
)

var (
	fConstants    bool
	fFields       bool
	fFunctions    bool
	fTypes        bool
	fVariables    bool
	fDebug        string
	fWholeProgram bool
	fReflection   bool
	fTags         string
)

func init() {
	flag.BoolVar(&fConstants, "consts", true, "Report unused constants")
	flag.BoolVar(&fFields, "fields", true, "Report unused fields")
	flag.BoolVar(&fFunctions, "funcs", true, "Report unused functions and methods")
	flag.BoolVar(&fTypes, "types", true, "Report unused types")
	flag.BoolVar(&fVariables, "vars", true, "Report unused variables")
	flag.StringVar(&fDebug, "debug", "", "Write a debug graph to `file`. Existing files will be overwritten.")
	flag.BoolVar(&fWholeProgram, "exported", false, "Treat arguments as a program and report unused exported identifiers")
	flag.BoolVar(&fReflection, "reflect", true, "Consider identifiers as used when it's likely they'll be accessed via reflection")
	flag.StringVar(&fTags, "tags", "", "Build tags")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] [packages]\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func newChecker(mode unused.CheckMode) *unused.Checker {
	checker := unused.NewChecker(mode)

	if fDebug != "" {
		debug, err := os.Create(fDebug)
		if err != nil {
			log.Fatal("couldn't open debug file:", err)
		}
		checker.Debug = debug
	}

	checker.WholeProgram = fWholeProgram
	checker.ConsiderReflection = fReflection
	return checker
}

func main() {
	log.SetFlags(0)
	flag.Parse()
	var mode unused.CheckMode
	if fConstants {
		mode |= unused.CheckConstants
	}
	if fFields {
		mode |= unused.CheckFields
	}
	if fFunctions {
		mode |= unused.CheckFunctions
	}
	if fTypes {
		mode |= unused.CheckTypes
	}
	if fVariables {
		mode |= unused.CheckVariables
	}

	checker := newChecker(mode)
	l := unused.NewLintChecker(checker)
	args := flag.Args()
	if len(fTags) > 0 {
		args = append([]string{"-tags", fTags}, args...)
	}
	lintutil.ProcessArgs("unused", l, args)
}
<commit_msg>Support lintutil's -ignore flag<commit_after>package main // import "honnef.co/go/unused/cmd/unused"

import (
	"flag"
	"fmt"
	"log"
	"os"

	"honnef.co/go/lint/lintutil"
	"honnef.co/go/unused"
)

var (
	fConstants    bool
	fFields       bool
	fFunctions    bool
	fTypes        bool
	fVariables    bool
	fDebug        string
	fWholeProgram bool
	fReflection   bool
	fTags         string
	fIgnore       string
)

func init() {
	flag.BoolVar(&fConstants, "consts", true, "Report unused constants")
	flag.BoolVar(&fFields, "fields", true, "Report unused fields")
	flag.BoolVar(&fFunctions, "funcs", true, "Report unused functions and methods")
	flag.BoolVar(&fTypes, "types", true, "Report unused types")
	flag.BoolVar(&fVariables, "vars", true, "Report unused variables")
	flag.StringVar(&fDebug, "debug", "", "Write a debug graph to `file`. Existing files will be overwritten.")
	flag.BoolVar(&fWholeProgram, "exported", false, "Treat arguments as a program and report unused exported identifiers")
	flag.BoolVar(&fReflection, "reflect", true, "Consider identifiers as used when it's likely they'll be accessed via reflection")
	flag.StringVar(&fTags, "tags", "", "List of `build tags`")
	flag.StringVar(&fIgnore, "ignore", "", "Space separated list of checks to ignore, in the following format: 'import/path/file.go:Check1,Check2,...' Both the import path and file name sections support globbing, e.g. 'os/exec/*_test.go'")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] [packages]\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func newChecker(mode unused.CheckMode) *unused.Checker {
	checker := unused.NewChecker(mode)

	if fDebug != "" {
		debug, err := os.Create(fDebug)
		if err != nil {
			log.Fatal("couldn't open debug file:", err)
		}
		checker.Debug = debug
	}

	checker.WholeProgram = fWholeProgram
	checker.ConsiderReflection = fReflection
	return checker
}

func main() {
	log.SetFlags(0)
	flag.Parse()
	var mode unused.CheckMode
	if fConstants {
		mode |= unused.CheckConstants
	}
	if fFields {
		mode |= unused.CheckFields
	}
	if fFunctions {
		mode |= unused.CheckFunctions
	}
	if fTypes {
		mode |= unused.CheckTypes
	}
	if fVariables {
		mode |= unused.CheckVariables
	}

	checker := newChecker(mode)
	l := unused.NewLintChecker(checker)
	args := []string{
		"-tags", fTags,
		"-ignore", fIgnore,
	}
	args = append(args, flag.Args()...)
	lintutil.ProcessArgs("unused", l, args)
}
<|endoftext|>


<commit_before><commit_msg>Reverting to mgoTimeout changed Dial() to DialWithTimeout() commit<commit_after><|endoftext|>


<commit_before>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package stop_test

import (
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	_ "github.com/cockroachdb/cockroach/util/log" // for flags
	"github.com/cockroachdb/cockroach/util/stop"
)

func TestStopper(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	running := make(chan struct{})
	waiting := make(chan struct{})
	cleanup := make(chan struct{})

	s.RunWorker(func() {
		<-running
	})

	go func() {
		<-s.ShouldStop()
		select {
		case <-waiting:
			t.Fatal("expected stopper to have blocked")
		case <-time.After(1 * time.Millisecond):
			// Expected.
		}
		close(running)
		select {
		case <-waiting:
			// Success.
		case <-time.After(100 * time.Millisecond):
			t.Fatal("stopper should have finished waiting")
		}
		close(cleanup)
	}()

	s.Stop()
	close(waiting)
	<-cleanup
}

type blockingCloser struct {
	block chan struct{}
}

func newBlockingCloser() *blockingCloser {
	return &blockingCloser{block: make(chan struct{})}
}

func (bc *blockingCloser) Unblock() {
	close(bc.block)
}

func (bc *blockingCloser) Close() {
	<-bc.block
}

func TestStopperIsStopped(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	bc := newBlockingCloser()
	s.AddCloser(bc)
	go s.Stop()

	select {
	case <-s.ShouldStop():
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should have finished waiting")
	}
	select {
	case <-s.IsStopped():
		t.Fatal("expected blocked closer to prevent stop")
	case <-time.After(1 * time.Millisecond):
		// Expected.
	}
	bc.Unblock()
	select {
	case <-s.IsStopped():
		// Expected
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should have finished stopping")
	}
}

func TestStopperMultipleStopees(t *testing.T) {
	defer leaktest.AfterTest(t)
	const count = 3
	s := stop.NewStopper()

	for i := 0; i < count; i++ {
		s.RunWorker(func() {
			<-s.ShouldStop()
		})
	}

	done := make(chan struct{})
	go func() {
		s.Stop()
		close(done)
	}()

	<-done
}

func TestStopperStartFinishTasks(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()

	if !s.RunTask(func() {
		go s.Stop()

		select {
		case <-s.ShouldStop():
			t.Fatal("expected stopper to be draining")
		case <-time.After(1 * time.Millisecond):
			// Expected.
		}
	}) {
		t.Error("expected RunTask to succeed")
	}
	select {
	case <-s.ShouldStop():
		// Success.
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should be ready to stop")
	}
}

func TestStopperRunWorker(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	s.RunWorker(func() {
		select {
		case <-s.ShouldStop():
			return
		}
	})
	closer := make(chan struct{})
	go func() {
		s.Stop()
		close(closer)
	}()
	select {
	case <-closer:
		// Success.
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should be ready to stop")
	}
}

// TestStopperQuiesce tests coordinate drain with Quiesce.
func TestStopperQuiesce(t *testing.T) {
	defer leaktest.AfterTest(t)
	var stoppers []*stop.Stopper
	for i := 0; i < 3; i++ {
		stoppers = append(stoppers, stop.NewStopper())
	}
	var quiesceDone []chan struct{}
	var runTaskDone []chan struct{}

	for _, s := range stoppers {
		thisStopper := s
		qc := make(chan struct{})
		quiesceDone = append(quiesceDone, qc)
		sc := make(chan struct{})
		runTaskDone = append(runTaskDone, sc)
		thisStopper.RunWorker(func() {
			// Wait until Quiesce() is called.
			<-qc
			if thisStopper.RunTask(func() {}) {
				t.Error("expected RunTask to fail")
			}
			// Make the stoppers call Stop().
			close(sc)
			<-thisStopper.ShouldStop()
		})
	}

	done := make(chan struct{})
	go func() {
		for _, s := range stoppers {
			s.Quiesce()
		}
		// Make the tasks call RunTask().
		for _, qc := range quiesceDone {
			close(qc)
		}

		// Wait until RunTask() is called.
		for _, sc := range runTaskDone {
			<-sc
		}

		for _, s := range stoppers {
			s.Stop()
			<-s.IsStopped()
		}
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(50 * time.Millisecond):
		t.Errorf("timed out waiting for stop")
	}
}

type testCloser bool

func (tc *testCloser) Close() {
	*tc = true
}

func TestStopperClosers(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	var tc1, tc2 testCloser
	s.AddCloser(&tc1)
	s.AddCloser(&tc2)
	s.Stop()
	if bool(tc1) != true || bool(tc2) != true {
		t.Errorf("expected true & true; got %t & %t", tc1, tc2)
	}
}

func TestStopperNumTasks(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	var tasks []chan bool
	for i := 0; i < 3; i++ {
		c := make(chan bool)
		tasks = append(tasks, c)
		s.RunAsyncTask(func() {
			// Wait for channel to close
			<-c
		})
		tm := s.RunningTasks()
		if numTypes, numTasks := len(tm), s.NumTasks(); numTypes != 1 || numTasks != i+1 {
			t.Errorf("stopper should have %d running tasks, got %d / %+v", i+1, numTasks, tm)
		}
		m := s.RunningTasks()
		if len(m) != 1 {
			t.Fatalf("expected exactly one task map entry: %+v", m)
		}
		for _, v := range m {
			if expNum := len(tasks); v != expNum {
				t.Fatalf("%d: expected %d tasks, got %d", i, expNum, v)
			}
		}
	}
	for i, c := range tasks {
		m := s.RunningTasks()
		if len(m) != 1 {
			t.Fatalf("%d: expected exactly one task map entry: %+v", i, m)
		}
		for _, v := range m {
			if expNum := len(tasks[i:]); v != expNum {
				t.Fatalf("%d: expected %d tasks, got %d:\n%s", i, expNum, v, m)
			}
		}
		// Close the channel to let the task proceed.
		close(c)
		expNum := len(tasks[i+1:])
		err := util.IsTrueWithin(func() bool { return s.NumTasks() == expNum }, 20*time.Millisecond)
		if err != nil {
			t.Errorf("%d: stopper should have %d running tasks, got %d", i, expNum, s.NumTasks())
		}
	}
	// The taskmap should've been cleared out.
	if m := s.RunningTasks(); len(m) != 0 {
		t.Fatalf("task map not empty: %+v", m)
	}
	s.Stop()
}

// TestStopperRunTaskPanic ensures that tasks are not leaked when they panic.
// RunAsyncTask has a similar bit of logic, but it is not testable because
// we cannot insert a recover() call in the right place.
func TestStopperRunTaskPanic(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	// If RunTask were not panic-safe, Stop() would deadlock.
	defer s.Stop()
	func() {
		defer func() {
			_ = recover()
		}()
		s.RunTask(func() {
			panic("ouch")
		})
	}()
}

func TestStopperShouldDrain(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	running := make(chan struct{})
	runningTask := make(chan struct{})
	waiting := make(chan struct{})
	cleanup := make(chan struct{})

	// Run a worker. A call to stopper.Stop() will not close until all workers
	// have completed, and this worker will complete when the "running" channel
	// is closed.
	s.RunWorker(func() {
		<-running
	})
	// Run an asynchronous task. A stopper which has been Stop()ed will not
	// close it's ShouldStop() channel until all tasks have completed. This task
	// will complete when the "runningTask" channel is closed.
	s.RunAsyncTask(func() {
		<-runningTask
	})

	go func() {
		// The ShouldDrain() channel should close as soon as the stopper is
		// Stop()ed.
		<-s.ShouldDrain()
		// However, the ShouldStop() channel should still be blocked because the
		// async task started above is still running, meaning we haven't drained
		// yet.
		select {
		case <-s.ShouldStop():
			t.Fatal("expected ShouldStop() to block until draining complete")
		default:
			// Expected.
		}
		// After completing the running task, the ShouldStop() channel should
		// now close.
		close(runningTask)
		<-s.ShouldStop()
		// However, the working running above prevents the call to Stop() from
		// returning; it blocks until the runner's goroutine is finished. We
		// use the "waiting" channel to detect this.
		select {
		case <-waiting:
			t.Fatal("expected stopper to have blocked")
		default:
			// Expected.
		}
		// Finally, close the "running" channel, which should cause the original
		// call to Stop() to return.
		close(running)
		<-waiting
		close(cleanup)
	}()

	s.Stop()
	close(waiting)
	<-cleanup
}
<commit_msg>add benchmark of util/stopper<commit_after>// Copyright 2014 The Cockroach 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.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

package stop_test

import (
	"fmt"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/leaktest"
	_ "github.com/cockroachdb/cockroach/util/log" // for flags
	"github.com/cockroachdb/cockroach/util/stop"
)

func TestStopper(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	running := make(chan struct{})
	waiting := make(chan struct{})
	cleanup := make(chan struct{})

	s.RunWorker(func() {
		<-running
	})

	go func() {
		<-s.ShouldStop()
		select {
		case <-waiting:
			t.Fatal("expected stopper to have blocked")
		case <-time.After(1 * time.Millisecond):
			// Expected.
		}
		close(running)
		select {
		case <-waiting:
			// Success.
		case <-time.After(100 * time.Millisecond):
			t.Fatal("stopper should have finished waiting")
		}
		close(cleanup)
	}()

	s.Stop()
	close(waiting)
	<-cleanup
}

type blockingCloser struct {
	block chan struct{}
}

func newBlockingCloser() *blockingCloser {
	return &blockingCloser{block: make(chan struct{})}
}

func (bc *blockingCloser) Unblock() {
	close(bc.block)
}

func (bc *blockingCloser) Close() {
	<-bc.block
}

func TestStopperIsStopped(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	bc := newBlockingCloser()
	s.AddCloser(bc)
	go s.Stop()

	select {
	case <-s.ShouldStop():
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should have finished waiting")
	}
	select {
	case <-s.IsStopped():
		t.Fatal("expected blocked closer to prevent stop")
	case <-time.After(1 * time.Millisecond):
		// Expected.
	}
	bc.Unblock()
	select {
	case <-s.IsStopped():
		// Expected
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should have finished stopping")
	}
}

func TestStopperMultipleStopees(t *testing.T) {
	defer leaktest.AfterTest(t)
	const count = 3
	s := stop.NewStopper()

	for i := 0; i < count; i++ {
		s.RunWorker(func() {
			<-s.ShouldStop()
		})
	}

	done := make(chan struct{})
	go func() {
		s.Stop()
		close(done)
	}()

	<-done
}

func TestStopperStartFinishTasks(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()

	if !s.RunTask(func() {
		go s.Stop()

		select {
		case <-s.ShouldStop():
			t.Fatal("expected stopper to be draining")
		case <-time.After(1 * time.Millisecond):
			// Expected.
		}
	}) {
		t.Error("expected RunTask to succeed")
	}
	select {
	case <-s.ShouldStop():
		// Success.
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should be ready to stop")
	}
}

func TestStopperRunWorker(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	s.RunWorker(func() {
		select {
		case <-s.ShouldStop():
			return
		}
	})
	closer := make(chan struct{})
	go func() {
		s.Stop()
		close(closer)
	}()
	select {
	case <-closer:
		// Success.
	case <-time.After(100 * time.Millisecond):
		t.Fatal("stopper should be ready to stop")
	}
}

// TestStopperQuiesce tests coordinate drain with Quiesce.
func TestStopperQuiesce(t *testing.T) {
	defer leaktest.AfterTest(t)
	var stoppers []*stop.Stopper
	for i := 0; i < 3; i++ {
		stoppers = append(stoppers, stop.NewStopper())
	}
	var quiesceDone []chan struct{}
	var runTaskDone []chan struct{}

	for _, s := range stoppers {
		thisStopper := s
		qc := make(chan struct{})
		quiesceDone = append(quiesceDone, qc)
		sc := make(chan struct{})
		runTaskDone = append(runTaskDone, sc)
		thisStopper.RunWorker(func() {
			// Wait until Quiesce() is called.
			<-qc
			if thisStopper.RunTask(func() {}) {
				t.Error("expected RunTask to fail")
			}
			// Make the stoppers call Stop().
			close(sc)
			<-thisStopper.ShouldStop()
		})
	}

	done := make(chan struct{})
	go func() {
		for _, s := range stoppers {
			s.Quiesce()
		}
		// Make the tasks call RunTask().
		for _, qc := range quiesceDone {
			close(qc)
		}

		// Wait until RunTask() is called.
		for _, sc := range runTaskDone {
			<-sc
		}

		for _, s := range stoppers {
			s.Stop()
			<-s.IsStopped()
		}
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(50 * time.Millisecond):
		t.Errorf("timed out waiting for stop")
	}
}

type testCloser bool

func (tc *testCloser) Close() {
	*tc = true
}

func TestStopperClosers(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	var tc1, tc2 testCloser
	s.AddCloser(&tc1)
	s.AddCloser(&tc2)
	s.Stop()
	if bool(tc1) != true || bool(tc2) != true {
		t.Errorf("expected true & true; got %t & %t", tc1, tc2)
	}
}

func TestStopperNumTasks(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	var tasks []chan bool
	for i := 0; i < 3; i++ {
		c := make(chan bool)
		tasks = append(tasks, c)
		s.RunAsyncTask(func() {
			// Wait for channel to close
			<-c
		})
		tm := s.RunningTasks()
		if numTypes, numTasks := len(tm), s.NumTasks(); numTypes != 1 || numTasks != i+1 {
			t.Errorf("stopper should have %d running tasks, got %d / %+v", i+1, numTasks, tm)
		}
		m := s.RunningTasks()
		if len(m) != 1 {
			t.Fatalf("expected exactly one task map entry: %+v", m)
		}
		for _, v := range m {
			if expNum := len(tasks); v != expNum {
				t.Fatalf("%d: expected %d tasks, got %d", i, expNum, v)
			}
		}
	}
	for i, c := range tasks {
		m := s.RunningTasks()
		if len(m) != 1 {
			t.Fatalf("%d: expected exactly one task map entry: %+v", i, m)
		}
		for _, v := range m {
			if expNum := len(tasks[i:]); v != expNum {
				t.Fatalf("%d: expected %d tasks, got %d:\n%s", i, expNum, v, m)
			}
		}
		// Close the channel to let the task proceed.
		close(c)
		expNum := len(tasks[i+1:])
		err := util.IsTrueWithin(func() bool { return s.NumTasks() == expNum }, 20*time.Millisecond)
		if err != nil {
			t.Errorf("%d: stopper should have %d running tasks, got %d", i, expNum, s.NumTasks())
		}
	}
	// The taskmap should've been cleared out.
	if m := s.RunningTasks(); len(m) != 0 {
		t.Fatalf("task map not empty: %+v", m)
	}
	s.Stop()
}

// TestStopperRunTaskPanic ensures that tasks are not leaked when they panic.
// RunAsyncTask has a similar bit of logic, but it is not testable because
// we cannot insert a recover() call in the right place.
func TestStopperRunTaskPanic(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	// If RunTask were not panic-safe, Stop() would deadlock.
	defer s.Stop()
	func() {
		defer func() {
			_ = recover()
		}()
		s.RunTask(func() {
			panic("ouch")
		})
	}()
}

func TestStopperShouldDrain(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := stop.NewStopper()
	running := make(chan struct{})
	runningTask := make(chan struct{})
	waiting := make(chan struct{})
	cleanup := make(chan struct{})

	// Run a worker. A call to stopper.Stop() will not close until all workers
	// have completed, and this worker will complete when the "running" channel
	// is closed.
	s.RunWorker(func() {
		<-running
	})
	// Run an asynchronous task. A stopper which has been Stop()ed will not
	// close it's ShouldStop() channel until all tasks have completed. This task
	// will complete when the "runningTask" channel is closed.
	s.RunAsyncTask(func() {
		<-runningTask
	})

	go func() {
		// The ShouldDrain() channel should close as soon as the stopper is
		// Stop()ed.
		<-s.ShouldDrain()
		// However, the ShouldStop() channel should still be blocked because the
		// async task started above is still running, meaning we haven't drained
		// yet.
		select {
		case <-s.ShouldStop():
			t.Fatal("expected ShouldStop() to block until draining complete")
		default:
			// Expected.
		}
		// After completing the running task, the ShouldStop() channel should
		// now close.
		close(runningTask)
		<-s.ShouldStop()
		// However, the working running above prevents the call to Stop() from
		// returning; it blocks until the runner's goroutine is finished. We
		// use the "waiting" channel to detect this.
		select {
		case <-waiting:
			t.Fatal("expected stopper to have blocked")
		default:
			// Expected.
		}
		// Finally, close the "running" channel, which should cause the original
		// call to Stop() to return.
		close(running)
		<-waiting
		close(cleanup)
	}()

	s.Stop()
	close(waiting)
	<-cleanup
}

func maybePrint() {
	if testing.Verbose() { // This just needs to be complicated enough not to inline.
		fmt.Println("blah")
	}
}

func BenchmarkDirectCall(b *testing.B) {
	defer leaktest.AfterTest(b)
	s := stop.NewStopper()
	defer s.Stop()
	for i := 0; i < b.N; i++ {
		maybePrint()
	}
}

func BenchmarkStopper(b *testing.B) {
	defer leaktest.AfterTest(b)
	s := stop.NewStopper()
	defer s.Stop()
	for i := 0; i < b.N; i++ {
		s.RunTask(maybePrint)
	}
}
func BenchmarkDirectCallPar(b *testing.B) {
	defer leaktest.AfterTest(b)
	s := stop.NewStopper()
	defer s.Stop()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			maybePrint()
		}
	})
}

func BenchmarkStopperPar(b *testing.B) {
	defer leaktest.AfterTest(b)
	s := stop.NewStopper()
	defer s.Stop()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			s.RunTask(maybePrint)
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2013 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.

// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.
// +build leak

package testleak

import (
	"runtime"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/pingcap/check"
)

func interestingGoroutines() (gs []string) {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	ignoreList := []string{
		"created by github.com/pingcap/tidb.init",
		"testing.RunTests",
		"check.(*resultTracker).start",
		"check.(*suiteRunner).runFunc",
		"check.(*suiteRunner).parallelRun",
		"localstore.(*dbStore).scheduler",
		"tikv.(*noGCHandler).Start",
		"ddl.(*ddl).start",
		"ddl.(*delRange).startEmulator",
		"domain.NewDomain",
		"testing.(*T).Run",
		"domain.(*Domain).LoadPrivilegeLoop",
		"domain.(*Domain).UpdateTableStatsLoop",
		"testing.Main(",
		"runtime.goexit",
		"created by runtime.gc",
		"interestingGoroutines",
		"runtime.MHeap_Scavenger",
		"created by os/signal.init",
		// these go routines are async terminated, so they may still alive after test end, thus cause
		// false positive leak failures
		"google.golang.org/grpc.(*addrConn).resetTransport",
		"google.golang.org/grpc.(*ccBalancerWrapper).watcher",
		"github.com/pingcap/goleveldb/leveldb/util.(*BufferPool).drain",
		"github.com/pingcap/goleveldb/leveldb.(*DB).compactionError",
		"github.com/pingcap/goleveldb/leveldb.(*DB).mpoolDrain",
		"go.etcd.io/etcd/v3/pkg/logutil.(*MergeLogger).outputLoop",
		"oracles.(*pdOracle).updateTS",
		"tikv.(*tikvStore).runSafePointChecker",
		"tikv.(*RegionCache).asyncCheckAndResolveLoop",
		"github.com/pingcap/badger.(*writeWorker).runMergeLSM",
	}
	shouldIgnore := func(stack string) bool {
		if stack == "" {
			return true
		}
		for _, ident := range ignoreList {
			if strings.Contains(stack, ident) {
				return true
			}
		}
		return false
	}
	for _, g := range strings.Split(string(buf), "\n\n") {
		sl := strings.SplitN(g, "\n", 2)
		if len(sl) != 2 {
			continue
		}
		stack := strings.TrimSpace(sl[1])
		if shouldIgnore(stack) {
			continue
		}
		gs = append(gs, stack)
	}
	sort.Strings(gs)
	return
}

var beforeTestGoroutines = map[string]bool{}
var testGoroutinesInited bool

// BeforeTest gets the current goroutines.
// It's used for check.Suite.SetUpSuite() function.
// Now it's only used in the tidb_test.go.
// Note: it's not accurate, consider the following function:
// func loop() {
//   for {
//     select {
//       case <-ticker.C:
//         DoSomething()
//     }
//   }
// }
// If this loop step into DoSomething() during BeforeTest(), the stack for this goroutine will contain DoSomething().
// Then if this loop jumps out of DoSomething during AfterTest(), the stack for this goroutine will not contain DoSomething().
// Resulting in false-positive leak reports.
func BeforeTest() {
	for _, g := range interestingGoroutines() {
		beforeTestGoroutines[g] = true
	}
	testGoroutinesInited = true
}

const defaultCheckCnt = 50

func checkLeakAfterTest(errorFunc func(cnt int, g string)) func() {
	// After `BeforeTest`, `beforeTestGoroutines` may still be empty, in this case,
	// we shouldn't init it again.
	if !testGoroutinesInited && len(beforeTestGoroutines) == 0 {
		for _, g := range interestingGoroutines() {
			beforeTestGoroutines[g] = true
		}
	}

	cnt := defaultCheckCnt
	return func() {
		defer func() {
			beforeTestGoroutines = map[string]bool{}
			testGoroutinesInited = false
		}()

		var leaked []string
		for i := 0; i < cnt; i++ {
			leaked = leaked[:0]
			for _, g := range interestingGoroutines() {
				if !beforeTestGoroutines[g] {
					leaked = append(leaked, g)
				}
			}
			// Bad stuff found, but goroutines might just still be
			// shutting down, so give it some time.
			if len(leaked) != 0 {
				time.Sleep(50 * time.Millisecond)
				continue
			}

			return
		}
		for _, g := range leaked {
			errorFunc(cnt, g)
		}
	}
}

// AfterTest gets the current goroutines and runs the returned function to
// get the goroutines at that time to contrast whether any goroutines leaked.
// Usage: defer testleak.AfterTest(c)()
// It can call with BeforeTest() at the beginning of check.Suite.TearDownSuite() or
// call alone at the beginning of each test.
func AfterTest(c *check.C) func() {
	errorFunc := func(cnt int, g string) {
		c.Errorf("Test %s check-count %d appears to have leaked: %v", c.TestName(), cnt, g)
	}
	return checkLeakAfterTest(errorFunc)
}

// AfterTestT is used after all the test cases is finished.
func AfterTestT(t *testing.T) func() {
	errorFunc := func(cnt int, g string) {
		t.Errorf("Test %s check-count %d appears to have leaked: %v", t.Name(), cnt, g)
	}
	return checkLeakAfterTest(errorFunc)
}
<commit_msg>test: stabilize leak test (#19526)<commit_after>// Copyright 2013 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.

// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.
// +build leak

package testleak

import (
	"runtime"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/pingcap/check"
)

func interestingGoroutines() (gs []string) {
	buf := make([]byte, 2<<20)
	buf = buf[:runtime.Stack(buf, true)]
	ignoreList := []string{
		"created by github.com/pingcap/tidb.init",
		"testing.RunTests",
		"check.(*resultTracker).start",
		"check.(*suiteRunner).runFunc",
		"check.(*suiteRunner).parallelRun",
		"localstore.(*dbStore).scheduler",
		"tikv.(*noGCHandler).Start",
		"ddl.(*ddl).start",
		"ddl.(*delRange).startEmulator",
		"domain.NewDomain",
		"testing.(*T).Run",
		"domain.(*Domain).LoadPrivilegeLoop",
		"domain.(*Domain).UpdateTableStatsLoop",
		"testing.Main(",
		"runtime.goexit",
		"created by runtime.gc",
		"interestingGoroutines",
		"runtime.MHeap_Scavenger",
		"created by os/signal.init",
		// these go routines are async terminated, so they may still alive after test end, thus cause
		// false positive leak failures
		"google.golang.org/grpc.(*addrConn).resetTransport",
		"google.golang.org/grpc.(*ccBalancerWrapper).watcher",
		"github.com/pingcap/goleveldb/leveldb/util.(*BufferPool).drain",
		"github.com/pingcap/goleveldb/leveldb.(*DB).compactionError",
		"github.com/pingcap/goleveldb/leveldb.(*DB).mpoolDrain",
		"go.etcd.io/etcd/pkg/logutil.(*MergeLogger).outputLoop",
		"go.etcd.io/etcd/v3/pkg/logutil.(*MergeLogger).outputLoop",
		"oracles.(*pdOracle).updateTS",
		"tikv.(*tikvStore).runSafePointChecker",
		"tikv.(*RegionCache).asyncCheckAndResolveLoop",
		"github.com/pingcap/badger",
		"github.com/ngaut/unistore/tikv.(*MVCCStore).runUpdateSafePointLoop",
	}
	shouldIgnore := func(stack string) bool {
		if stack == "" {
			return true
		}
		for _, ident := range ignoreList {
			if strings.Contains(stack, ident) {
				return true
			}
		}
		return false
	}
	for _, g := range strings.Split(string(buf), "\n\n") {
		sl := strings.SplitN(g, "\n", 2)
		if len(sl) != 2 {
			continue
		}
		stack := strings.TrimSpace(sl[1])
		if shouldIgnore(stack) {
			continue
		}
		gs = append(gs, stack)
	}
	sort.Strings(gs)
	return
}

var beforeTestGoroutines = map[string]bool{}
var testGoroutinesInited bool

// BeforeTest gets the current goroutines.
// It's used for check.Suite.SetUpSuite() function.
// Now it's only used in the tidb_test.go.
// Note: it's not accurate, consider the following function:
// func loop() {
//   for {
//     select {
//       case <-ticker.C:
//         DoSomething()
//     }
//   }
// }
// If this loop step into DoSomething() during BeforeTest(), the stack for this goroutine will contain DoSomething().
// Then if this loop jumps out of DoSomething during AfterTest(), the stack for this goroutine will not contain DoSomething().
// Resulting in false-positive leak reports.
func BeforeTest() {
	for _, g := range interestingGoroutines() {
		beforeTestGoroutines[g] = true
	}
	testGoroutinesInited = true
}

const defaultCheckCnt = 50

func checkLeakAfterTest(errorFunc func(cnt int, g string)) func() {
	// After `BeforeTest`, `beforeTestGoroutines` may still be empty, in this case,
	// we shouldn't init it again.
	if !testGoroutinesInited && len(beforeTestGoroutines) == 0 {
		for _, g := range interestingGoroutines() {
			beforeTestGoroutines[g] = true
		}
	}

	cnt := defaultCheckCnt
	return func() {
		defer func() {
			beforeTestGoroutines = map[string]bool{}
			testGoroutinesInited = false
		}()

		var leaked []string
		for i := 0; i < cnt; i++ {
			leaked = leaked[:0]
			for _, g := range interestingGoroutines() {
				if !beforeTestGoroutines[g] {
					leaked = append(leaked, g)
				}
			}
			// Bad stuff found, but goroutines might just still be
			// shutting down, so give it some time.
			if len(leaked) != 0 {
				time.Sleep(50 * time.Millisecond)
				continue
			}

			return
		}
		for _, g := range leaked {
			errorFunc(cnt, g)
		}
	}
}

// AfterTest gets the current goroutines and runs the returned function to
// get the goroutines at that time to contrast whether any goroutines leaked.
// Usage: defer testleak.AfterTest(c)()
// It can call with BeforeTest() at the beginning of check.Suite.TearDownSuite() or
// call alone at the beginning of each test.
func AfterTest(c *check.C) func() {
	errorFunc := func(cnt int, g string) {
		c.Errorf("Test %s check-count %d appears to have leaked: %v", c.TestName(), cnt, g)
	}
	return checkLeakAfterTest(errorFunc)
}

// AfterTestT is used after all the test cases is finished.
func AfterTestT(t *testing.T) func() {
	errorFunc := func(cnt int, g string) {
		t.Errorf("Test %s check-count %d appears to have leaked: %v", t.Name(), cnt, g)
	}
	return checkLeakAfterTest(errorFunc)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
)

func main() {
	// for who.go Who(char, lvl)
	var char = flag.String("char", "", "Character name for update or import. Ex: Rynshana")
	var lvl = flag.Int("lvl", 0, "Character level for update or import. Ex: 50")
	// for who.go WhoChar(char, lvl, class, race, acct)
	var class = flag.String("class", "", "Character class for initial import. Ex: \"Cleric\"")
	var race = flag.String("race", "", "Character race for initial import. Ex: \"Moon Elf\"")
	var acct = flag.String("acct", "", "Character account for initial import. Ex: Krimic")
	// for identify.go Identify(filename)
	var file = flag.String("import", "", "Parse file for identify stats, import to DB. Ex: newstats.txt")
	// for time.go Uptime(curup)
	var time = flag.String("time", "", "Parse uptime for boot tracking. Ex: 58:10:26")
	// for tell.go
	var cmd = flag.String("cmd", "", "Command from tell. Ex: stat")
	var oper = flag.String("oper", "", "Operant from tell, to be operated on by cmd. Ex: \"a longsword\"")

	flag.Parse()
	*cmd += *oper

	// only run one command at a time
	if *char != "" && *lvl > 0 && *class != "" && *race != "" && *acct != "" {
		WhoChar(*char, *lvl, *class, *race, *acct)
	} else if *char != "" && 50 >= *lvl && *lvl > 0 {
		Who(*char, *lvl)
	} else if *file != "" {
		Identify(*file)
	} else if *time != "" {
		Uptime(*time)
	}
}
<commit_msg>More explicit check to make sure lvl makes sense<commit_after>package main

import (
	"flag"
)

func main() {
	// for who.go Who(char, lvl)
	var char = flag.String("char", "", "Character name for update or import. Ex: Rynshana")
	var lvl = flag.Int("lvl", 0, "Character level for update or import. Ex: 50")
	// for who.go WhoChar(char, lvl, class, race, acct)
	var class = flag.String("class", "", "Character class for initial import. Ex: \"Cleric\"")
	var race = flag.String("race", "", "Character race for initial import. Ex: \"Moon Elf\"")
	var acct = flag.String("acct", "", "Character account for initial import. Ex: Krimic")
	// for identify.go Identify(filename)
	var file = flag.String("import", "", "Parse file for identify stats, import to DB. Ex: newstats.txt")
	// for time.go Uptime(curup)
	var time = flag.String("time", "", "Parse uptime for boot tracking. Ex: 58:10:26")
	// for tell.go
	var cmd = flag.String("cmd", "", "Command from tell. Ex: stat")
	var oper = flag.String("oper", "", "Operant from tell, to be operated on by cmd. Ex: \"a longsword\"")

	flag.Parse()
	*cmd += *oper

	// only run one command at a time
	if *char != "" && 50 >= *lvl && *lvl > 0 && *class != "" && *race != "" && *acct != "" {
		WhoChar(*char, *lvl, *class, *race, *acct)
	} else if *char != "" && 50 >= *lvl && *lvl > 0 {
		Who(*char, *lvl)
	} else if *file != "" {
		Identify(*file)
	} else if *time != "" {
		Uptime(*time)
	}
}
<|endoftext|>