<commit_before>package oracle

import (
	"fmt"
	"testing"
	"time"
)

var dataTypesTests = []struct {
	in  string
	out string
}{
	{"SELECT 1 FROM DUAL", "%!s(float64=1)"},
	{"SELECT -1 FROM DUAL", "%!s(float64=-1)"},
	{"SELECT -1/4 FROM DUAL", "%!s(float64=-0.25)"},
	{"SELECT TO_DATE('2011-12-13 14:15:16', 'YYYY-MM-DD HH24:MI:SS') FROM DUAL",
		"2011-12-13 14:15:16 +0100 CET"},
	{"SELECT 'AbraKA' FROM DUAL", "AbraKA"},
	{"SELECT 'árvíztűrő tükörfúrógép' FROM DUAL", "árvíztűrő tükörfúrógép"},
	{"SELECT HEXTORAW('00') FROM DUAL", "\x00"},
	// {"SELECT TO_CLOB('árvíztűrő tükörfúrógép') FROM DUAL", "árvíztűrő tükörfúrógép"},
}

func TestSimpleTypes(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	var (
		row []interface{}
	)
	cur := conn.NewCursor()
	defer cur.Close()
	{
		if oci, client, db, err := conn.NlsSettings(cur); err != nil {
			t.Logf("NlsSettings: %s", err)
		} else {
			t.Logf("NLS oci=%s client=%s database=%s", oci, client, db)
		}
	}

	var (
		repr string
		err  error
	)
	for i, tt := range dataTypesTests {
		if err = cur.Execute(tt.in, nil, nil); err != nil {
			t.Errorf("error executing `%s`: %s", tt.in, err)
		} else {
			if row, err = cur.FetchOne(); err != nil {
				t.Errorf("cannot fetch row: %s", err)
			} else {
				repr = fmt.Sprintf("%s", row[0])
				if repr != tt.out {
					t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, repr, tt.out)
				}
			}
		}
	}
}

var bindsTests = []struct {
	in  interface{}
	out string
}{
	{1, "Typ=2 Len=2: c1,2"},
	{1.0 / 2, "Typ=2 Len=2: c0,33"},
	{-1.25, "Typ=2 Len=4: 3e,64,4c,66"},
	{"SELECT", "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	{"árvíztűrő tükörfúrógép", "Typ=1 Len=31 CharacterSet=AL32UTF8: c3,a1,72,76,c3,ad,7a,74,c5,b1,72,c5,91,20,74,c3,bc,6b,c3,b6,72,66,c3,ba,72,c3,b3,67,c3,a9,70"},
	{[]byte{0, 1, 2, 3, 5, 7, 11, 13}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
		"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestSimpleBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err  error
		row  []interface{}
		repr string
	)
	for i, tt := range bindsTests {
		if err = cur.Execute("SELECT DUMP(:1, 1016) FROM DUAL", []interface{}{tt.in}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", tt.in, err)
		} else {
			if row, err = cur.FetchOne(); err != nil {
				t.Errorf("cannot fetch row: %s", err)
			} else {
				repr = fmt.Sprintf("%s", row[0])
				if repr != tt.out {
					t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, repr, tt.out)
				}
			}
		}
	}
}

var outBindsTests = []struct {
	in_str  string
	out_val interface{}
	out_str string
}{
	{"3", int32(0), "3"},
	{"-10.24", float32(-10.24), "-10.24"},
	// {"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	// {"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
	// 	time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
	// 	"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestOutBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err          error
		qry, out_str string
		out          *Variable
	)
	for i, tt := range outBindsTests {
		qry = `BEGIN SELECT ` + tt.in_str + ` INTO :1 FROM DUAL; END;`
		if out, err = cur.NewVar(tt.out_val); err != nil {
			t.Errorf("error creating variable for %s(%T): %s", tt.out_val, tt.out_val, err)
		}
		if err = cur.Execute(qry, []interface{}{out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if err = out.GetValueInto(&tt.out_val, 0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		t.Logf("%d. out:%s %v", i, out, tt.out_val)
		out_str = fmt.Sprintf("%v", tt.out_val)
		if out_str != tt.out_str {
			t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in_str, out_str, tt.out_str)
		}
	}
}

var inOutBindsTests = []struct {
	in_typ string
	in     interface{}
	out    string
}{
	{"INTEGER(3)", int32(1), "Typ=2 Len=2: 193,2"},
	// {"NUMBER(5,3)", []float32{1.0 / 2, -10.24}, "Typ=2 Len=2: c0,33"},
	// {"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	// {"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
	// 	time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
	// 	"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestInOutBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		ok      bool
		err     error
		qry     string
		out     *Variable
		val     interface{}
		out_str string
	)
	if out, err = cur.NewVar(""); err != nil {
		t.Errorf("error creating output variable: %s", err)
		t.FailNow()
	}

	for i, tt := range inOutBindsTests {
		qry = `DECLARE
	v_in ` + tt.in_typ + ` := :1;
	v_out VARCHAR2(1000);
BEGIN
	SELECT DUMP(v_in) INTO v_out FROM DUAL;
	:2 := v_out;
END;`
		if err = cur.Execute(qry, []interface{}{tt.in, out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if val, err = out.GetValue(0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		if out_str, ok = val.(string); !ok {
			t.Logf("output is not string!?!, but %T (%v)", val, val)
		}
		t.Logf("%d. out:%s %v", i, out_str, val)
		if out_str != tt.out {
			t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, out_str, tt.out)
		}
	}
}

var arrBindsTests = []struct {
	tab_typ string
	in      interface{}
	out     string
}{
	{"INTEGER(3)", []int32{1, 3}, "Typ=2 Len=2: c1,2"},
	// {"NUMBER(5,3)", []float32{1.0 / 2, -10.24}, "Typ=2 Len=2: c0,33"},
	// {"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	// {"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
	// 	time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
	// 	"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestArrayBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err error
		qry string
		out *Variable
		val interface{}
	)
	for i, tt := range arrBindsTests {
		qry = `DECLARE
	tab_typ TABLE OF ` + tt.tab_typ + ` INDEX BY PLS_INTEGER;
	tab tab_typ := :1;
	v_idx PLS_INTEGER;
	out VARCHAR2(1000);
BEGIN
	v_idx := tab.FIRST;
	IF v_idx IS NULL THEN
		out := 'EMPTY';
	ELSE
	END IF;
	WHILE v_idx IS NOT NULL LOOP
		out := out||v_idx||'. '||DUMP(tab(v_idx))||CHR(10);
		v_idx := tab.NEXT(v_idx);
	END LOOP;
	:2 := out;
END;`
		if err = cur.Execute(qry, []interface{}{tt.in, out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if val, err = out.GetValue(0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		t.Logf("%d. out:%s %v", i, out, val)
		// if out != tt.out {
		// 	t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, out, tt.out)
		// }
	}
}
<commit_msg>arrays cannot get in...<commit_after>package oracle

import (
	"fmt"
	"testing"
	"time"
)

var dataTypesTests = []struct {
	in  string
	out string
}{
	{"SELECT 1 FROM DUAL", "%!s(float64=1)"},
	{"SELECT -1 FROM DUAL", "%!s(float64=-1)"},
	{"SELECT -1/4 FROM DUAL", "%!s(float64=-0.25)"},
	{"SELECT TO_DATE('2011-12-13 14:15:16', 'YYYY-MM-DD HH24:MI:SS') FROM DUAL",
		"2011-12-13 14:15:16 +0100 CET"},
	{"SELECT 'AbraKA' FROM DUAL", "AbraKA"},
	{"SELECT 'árvíztűrő tükörfúrógép' FROM DUAL", "árvíztűrő tükörfúrógép"},
	{"SELECT HEXTORAW('00') FROM DUAL", "\x00"},
	// {"SELECT TO_CLOB('árvíztűrő tükörfúrógép') FROM DUAL", "árvíztűrő tükörfúrógép"},
}

func TestSimpleTypes(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	var (
		row []interface{}
	)
	cur := conn.NewCursor()
	defer cur.Close()
	{
		if oci, client, db, err := conn.NlsSettings(cur); err != nil {
			t.Logf("NlsSettings: %s", err)
		} else {
			t.Logf("NLS oci=%s client=%s database=%s", oci, client, db)
		}
	}

	var (
		repr string
		err  error
	)
	for i, tt := range dataTypesTests {
		if err = cur.Execute(tt.in, nil, nil); err != nil {
			t.Errorf("error executing `%s`: %s", tt.in, err)
		} else {
			if row, err = cur.FetchOne(); err != nil {
				t.Errorf("cannot fetch row: %s", err)
			} else {
				repr = fmt.Sprintf("%s", row[0])
				if repr != tt.out {
					t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, repr, tt.out)
				}
			}
		}
	}
}

var bindsTests = []struct {
	in  interface{}
	out string
}{
	{1, "Typ=2 Len=2: c1,2"},
	{1.0 / 2, "Typ=2 Len=2: c0,33"},
	{-1.25, "Typ=2 Len=4: 3e,64,4c,66"},
	{"SELECT", "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	{"árvíztűrő tükörfúrógép", "Typ=1 Len=31 CharacterSet=AL32UTF8: c3,a1,72,76,c3,ad,7a,74,c5,b1,72,c5,91,20,74,c3,bc,6b,c3,b6,72,66,c3,ba,72,c3,b3,67,c3,a9,70"},
	{[]byte{0, 1, 2, 3, 5, 7, 11, 13}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
		"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestSimpleBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err  error
		row  []interface{}
		repr string
	)
	for i, tt := range bindsTests {
		if err = cur.Execute("SELECT DUMP(:1, 1016) FROM DUAL", []interface{}{tt.in}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", tt.in, err)
		} else {
			if row, err = cur.FetchOne(); err != nil {
				t.Errorf("cannot fetch row: %s", err)
			} else {
				repr = fmt.Sprintf("%s", row[0])
				if repr != tt.out {
					t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, repr, tt.out)
				}
			}
		}
	}
}

var outBindsTests = []struct {
	in_str  string
	out_val interface{}
	out_str string
}{
	{"3", int32(0), "3"},
	{"-10.24", float32(-10.24), "-10.24"},
	// {"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	// {"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
	// 	time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
	// 	"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestOutBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err          error
		qry, out_str string
		out          *Variable
	)
	for i, tt := range outBindsTests {
		qry = `BEGIN SELECT ` + tt.in_str + ` INTO :1 FROM DUAL; END;`
		if out, err = cur.NewVar(tt.out_val); err != nil {
			t.Errorf("error creating variable for %s(%T): %s", tt.out_val, tt.out_val, err)
		}
		if err = cur.Execute(qry, []interface{}{out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if err = out.GetValueInto(&tt.out_val, 0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		t.Logf("%d. out:%s %v", i, out, tt.out_val)
		out_str = fmt.Sprintf("%v", tt.out_val)
		if out_str != tt.out_str {
			t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in_str, out_str, tt.out_str)
		}
	}
}

var inOutBindsTests = []struct {
	in_typ string
	in     interface{}
	out    string
}{
	{"INTEGER(3)", int32(1), "Typ=2 Len=2: 193,2"},
	{"NUMBER(5,3)", []float32{1.0 / 2, -10.24}, "Typ=2 Len=2: c0,33"},
	{"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	{"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
		time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
		"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestInOutBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		ok      bool
		err     error
		qry     string
		out     *Variable
		val     interface{}
		out_str string
	)
	if out, err = cur.NewVar(""); err != nil {
		t.Errorf("error creating output variable: %s", err)
		t.FailNow()
	}

	for i, tt := range inOutBindsTests {
		qry = `DECLARE
	v_in ` + tt.in_typ + ` := :1;
	v_out VARCHAR2(1000);
BEGIN
	SELECT DUMP(v_in) INTO v_out FROM DUAL;
	:2 := v_out;
END;`
		if err = cur.Execute(qry, []interface{}{tt.in, out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if val, err = out.GetValue(0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		if out_str, ok = val.(string); !ok {
			t.Logf("output is not string!?!, but %T (%v)", val, val)
		}
		t.Logf("%d. out:%s %v", i, out_str, val)
		if out_str != tt.out {
			t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, out_str, tt.out)
		}
	}
}

var arrBindsTests = []struct {
	tab_typ string
	in      interface{}
	out     string
}{
	{"INTEGER(3)", []int32{1, 3, 5}, "Typ=2 Len=2: c1,2"},
	// {"NUMBER(5,3)", []float32{1.0 / 2, -10.24}, "Typ=2 Len=2: c0,33"},
	// {"VARCHAR2(40)", []string{"SELECT", "árvíztűrő tükörfúrógép"}, "Typ=1 Len=6 CharacterSet=AL32UTF8: 53,45,4c,45,43,54"},
	// {"RAW(4)", [][]byte{[]byte{0, 1, 2, 3}, []byte{5, 7, 11, 13}}, "Typ=23 Len=8: 0,1,2,3,5,7,b,d"},
	// {"DATE", []time.Time{time.Date(2013, 1, 2, 10, 6, 49, 0, time.Local),
	// 	time.Date(2012, 1, 2, 10, 6, 49, 0, time.Local)},
	// 	"Typ=12 Len=7: 78,71,1,2,b,7,32"},
}

func TestArrayBinds(t *testing.T) {
	conn := getConnection(t)
	if !conn.IsConnected() {
		t.FailNow()
	}
	cur := conn.NewCursor()
	defer cur.Close()

	var (
		err error
		qry string
		out *Variable
		val interface{}
	)
	if out, err = cur.NewVar(""); err != nil {
		t.Errorf("cannot create out variable: %s", err)
		t.FailNow()
	}
	for i, tt := range arrBindsTests {
		qry = `DECLARE
	TYPE tab_typ IS TABLE OF ` + tt.tab_typ + ` INDEX BY PLS_INTEGER;
	tab tab_typ;
	v_idx PLS_INTEGER;
	v_out VARCHAR2(1000);
BEGIN
	SELECT DUMP(:1) INTO v_out FROM DUAL;
	--tab := :1;
	v_idx := tab.FIRST;
	IF FALSE and v_idx IS NULL THEN
		v_out := 'EMPTY';
	END IF;
	WHILE v_idx IS NOT NULL LOOP
	    SELECT v_out||v_idx||'. '||DUMP(tab(v_idx))||CHR(10) INTO v_out FROM DUAL;
		v_idx := tab.NEXT(v_idx);
	END LOOP;
	:2 := v_out;
END;`
		if err = cur.Execute(qry, []interface{}{tt.in, out}, nil); err != nil {
			t.Errorf("error executing `%s`: %s", qry, err)
			continue
		}
		if val, err = out.GetValue(0); err != nil {
			t.Errorf("%d. error getting value: %s", i, err)
			continue
		}
		t.Logf("%d. out:%s => %v", i, out, val)
		// if out != tt.out {
		// 	t.Errorf("%d. exec(%q) => %q, want %q", i, tt.in, out, tt.out)
		// }
	}
}
<|endoftext|>


<commit_before>package pnet

import (
	"crypto/cipher"
	"crypto/rand"
	"io"
	"net"

	salsa20 "github.com/davidlazar/go-crypto/salsa20"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	mpool "github.com/libp2p/go-msgio/mpool"
)

// we are using buffer pool as user needs their slice back
// so we can't do XOR cripter in place
var (
	errShortNonce  = ipnet.NewError("could not read full nonce")
	errInsecureNil = ipnet.NewError("insecure is nil")
	errPSKNil      = ipnet.NewError("pre-shread key is nil")
)

type pskConn struct {
	net.Conn
	psk *[32]byte

	writeS20 cipher.Stream
	readS20  cipher.Stream
}

func (c *pskConn) Read(out []byte) (int, error) {
	if c.readS20 == nil {
		nonce := make([]byte, 24)
		_, err := io.ReadFull(c.Conn, nonce)
		if err != nil {
			return 0, errShortNonce
		}
		c.readS20 = salsa20.New(c.psk, nonce)
	}

	maxn := uint32(len(out))
	in := mpool.ByteSlicePool.Get(maxn).([]byte) // get buffer
	defer mpool.ByteSlicePool.Put(maxn, in)      // put the buffer back

	in = in[:maxn]            // truncate to required length
	n, err := c.Conn.Read(in) // read to in
	if n > 0 {
		c.readS20.XORKeyStream(out[:n], in[:n]) // decrypt to out buffer
	}
	return n, err
}

func (c *pskConn) Write(in []byte) (int, error) {
	if c.writeS20 == nil {
		nonce := make([]byte, 24)
		_, err := rand.Read(nonce)
		if err != nil {
			return 0, err
		}
		_, err = c.Conn.Write(nonce)
		if err != nil {
			return 0, err
		}

		c.writeS20 = salsa20.New(c.psk, nonce)
	}
	n := uint32(len(in))
	out := mpool.ByteSlicePool.Get(n).([]byte) // get buffer
	defer mpool.ByteSlicePool.Put(n, out)      // put the buffer back

	out = out[:n]                    // truncate to required length
	c.writeS20.XORKeyStream(out, in) // encrypt

	return c.Conn.Write(out) // send
}

var _ net.Conn = (*pskConn)(nil)

func newPSKConn(psk *[32]byte, insecure net.Conn) (net.Conn, error) {
	if insecure == nil {
		return nil, errInsecureNil
	}
	if psk == nil {
		return nil, errPSKNil
	}
	return &pskConn{
		Conn: insecure,
		psk:  psk,
	}, nil
}
<commit_msg>avoid the buffer-pool on read<commit_after>package pnet

import (
	"crypto/cipher"
	"crypto/rand"
	"io"
	"net"

	salsa20 "github.com/davidlazar/go-crypto/salsa20"
	ipnet "github.com/libp2p/go-libp2p-interface-pnet"
	mpool "github.com/libp2p/go-msgio/mpool"
)

// we are using buffer pool as user needs their slice back
// so we can't do XOR cripter in place
var (
	errShortNonce  = ipnet.NewError("could not read full nonce")
	errInsecureNil = ipnet.NewError("insecure is nil")
	errPSKNil      = ipnet.NewError("pre-shread key is nil")
)

type pskConn struct {
	net.Conn
	psk *[32]byte

	writeS20 cipher.Stream
	readS20  cipher.Stream
}

func (c *pskConn) Read(out []byte) (int, error) {
	if c.readS20 == nil {
		nonce := make([]byte, 24)
		_, err := io.ReadFull(c.Conn, nonce)
		if err != nil {
			return 0, errShortNonce
		}
		c.readS20 = salsa20.New(c.psk, nonce)
	}

	n, err := c.Conn.Read(out) // read to in
	if n > 0 {
		c.readS20.XORKeyStream(out[:n], out[:n]) // decrypt to out buffer
	}
	return n, err
}

func (c *pskConn) Write(in []byte) (int, error) {
	if c.writeS20 == nil {
		nonce := make([]byte, 24)
		_, err := rand.Read(nonce)
		if err != nil {
			return 0, err
		}
		_, err = c.Conn.Write(nonce)
		if err != nil {
			return 0, err
		}

		c.writeS20 = salsa20.New(c.psk, nonce)
	}
	n := uint32(len(in))
	out := mpool.ByteSlicePool.Get(n).([]byte) // get buffer
	defer mpool.ByteSlicePool.Put(n, out)      // put the buffer back

	out = out[:n]                    // truncate to required length
	c.writeS20.XORKeyStream(out, in) // encrypt

	return c.Conn.Write(out) // send
}

var _ net.Conn = (*pskConn)(nil)

func newPSKConn(psk *[32]byte, insecure net.Conn) (net.Conn, error) {
	if insecure == nil {
		return nil, errInsecureNil
	}
	if psk == nil {
		return nil, errPSKNil
	}
	return &pskConn{
		Conn: insecure,
		psk:  psk,
	}, nil
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
			pattern:   pattern,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<commit_msg>fix panic when reuseport was no available<commit_after>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
			pattern:   pattern,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		pattern:   pattern,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	lgbl "github.com/ipfs/go-libp2p-loggables"
	logging "github.com/ipfs/go-log"
	ma "github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-multiaddr-net"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	logdial := lgbl.Dial("conn", "", "", d.laddr, raddr)
	rpev := log.EventBegin(ctx, "tptDialReusePort", logdial)

	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	_ = ctx // TODO: implement DialContext in reuseport
	con, err := d.rd.Dial(network, netraddr)
	if err == nil {
		logdial["reuseport"] = "success"
		rpev.Done()
		return manet.WrapNetConn(con)
	}

	if !ReuseErrShouldRetry(err) {
		logdial["reuseport"] = "failure"
		logdial["error"] = err
		rpev.Done()
		return nil, err
	}

	logdial["reuseport"] = "retry"
	logdial["error"] = err
	rpev.Done()

	return d.madialer.DialContext(ctx, raddr)
}

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<commit_msg>fix: use DialContext of go-reuseport<commit_after>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	lgbl "github.com/ipfs/go-libp2p-loggables"
	logging "github.com/ipfs/go-log"
	ma "github.com/jbenet/go-multiaddr"
	manet "github.com/jbenet/go-multiaddr-net"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	logdial := lgbl.Dial("conn", "", "", d.laddr, raddr)
	rpev := log.EventBegin(ctx, "tptDialReusePort", logdial)

	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	_ = ctx // TODO: implement DialContext in reuseport
	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		logdial["reuseport"] = "success"
		rpev.Done()
		return manet.WrapNetConn(con)
	}

	if !ReuseErrShouldRetry(err) {
		logdial["reuseport"] = "failure"
		logdial["error"] = err
		rpev.Done()
		return nil, err
	}

	logdial["reuseport"] = "retry"
	logdial["error"] = err
	rpev.Done()

	return d.madialer.DialContext(ctx, raddr)
}

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	// get the local net.Addr manually
	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}

	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: net.Dialer{
				LocalAddr: la,
				Timeout:   base.Timeout,
			},
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
			pattern:   pattern,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		pattern:   pattern,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<commit_msg>fix source address not being set in non REUSEPORT dialer<commit_after>package tcp

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	reuseport "github.com/jbenet/go-reuseport"
	tpt "github.com/libp2p/go-libp2p-transport"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
	mafmt "github.com/whyrusleeping/mafmt"
)

var log = logging.Logger("tcp-tpt")

type TcpTransport struct {
	dlock   sync.Mutex
	dialers map[string]tpt.Dialer

	llock     sync.Mutex
	listeners map[string]tpt.Listener
}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport() *TcpTransport {
	return &TcpTransport{
		dialers:   make(map[string]tpt.Dialer),
		listeners: make(map[string]tpt.Listener),
	}
}

func (t *TcpTransport) Dialer(laddr ma.Multiaddr, opts ...tpt.DialOpt) (tpt.Dialer, error) {
	if laddr == nil {
		zaddr, err := ma.NewMultiaddr("/ip4/0.0.0.0/tcp/0")
		if err != nil {
			return nil, err
		}
		laddr = zaddr
	}
	t.dlock.Lock()
	defer t.dlock.Unlock()
	s := laddr.String()
	d, found := t.dialers[s]
	if found {
		return d, nil
	}
	var base manet.Dialer

	la, err := manet.ToNetAddr(laddr)
	if err != nil {
		return nil, err // something wrong with laddr.
	}
	base.Dialer.LocalAddr = la

	var doReuse bool
	for _, o := range opts {
		switch o := o.(type) {
		case tpt.TimeoutOpt:
			base.Timeout = time.Duration(o)
		case tpt.ReuseportOpt:
			doReuse = bool(o)
		default:
			return nil, fmt.Errorf("unrecognized option: %#v", o)
		}
	}

	tcpd, err := t.newTcpDialer(base, laddr, doReuse)
	if err != nil {
		return nil, err
	}

	t.dialers[s] = tcpd
	return tcpd, nil
}

func (t *TcpTransport) Listen(laddr ma.Multiaddr) (tpt.Listener, error) {
	if !t.Matches(laddr) {
		return nil, fmt.Errorf("tcp transport cannot listen on %q", laddr)
	}

	t.llock.Lock()
	defer t.llock.Unlock()
	s := laddr.String()
	l, found := t.listeners[s]
	if found {
		return l, nil
	}

	list, err := manetListen(laddr)
	if err != nil {
		return nil, err
	}

	tlist := &tcpListener{
		list:      list,
		transport: t,
	}

	t.listeners[s] = tlist
	return tlist, nil
}

func manetListen(addr ma.Multiaddr) (manet.Listener, error) {
	network, naddr, err := manet.DialArgs(addr)
	if err != nil {
		return nil, err
	}

	if ReuseportIsAvailable() {
		nl, err := reuseport.Listen(network, naddr)
		if err == nil {
			// hey, it worked!
			return manet.WrapNetListener(nl)
		}
		// reuseport is available, but we failed to listen. log debug, and retry normally.
		log.Debugf("reuseport available, but failed to listen: %s %s, %s", network, naddr, err)
	}

	// either reuseport not available, or it failed. try normally.
	return manet.Listen(addr)
}

func (t *TcpTransport) Matches(a ma.Multiaddr) bool {
	return mafmt.TCP.Matches(a)
}

type tcpDialer struct {
	laddr ma.Multiaddr

	doReuse bool

	rd       reuseport.Dialer
	madialer manet.Dialer
	pattern  mafmt.Pattern

	transport tpt.Transport
}

func (t *TcpTransport) newTcpDialer(base manet.Dialer, laddr ma.Multiaddr, doReuse bool) (*tcpDialer, error) {
	var pattern mafmt.Pattern
	if TCP4.Matches(laddr) {
		pattern = TCP4
	} else if TCP6.Matches(laddr) {
		pattern = TCP6
	} else {
		return nil, fmt.Errorf("local addr did not match TCP4 or TCP6: %s", laddr)
	}

	if doReuse && ReuseportIsAvailable() {
		rd := reuseport.Dialer{
			D: base.Dialer,
		}

		return &tcpDialer{
			doReuse:   true,
			laddr:     laddr,
			rd:        rd,
			madialer:  base,
			transport: t,
			pattern:   pattern,
		}, nil
	}

	return &tcpDialer{
		doReuse:   false,
		laddr:     laddr,
		pattern:   pattern,
		madialer:  base,
		transport: t,
	}, nil
}

func (d *tcpDialer) Dial(raddr ma.Multiaddr) (tpt.Conn, error) {
	return d.DialContext(context.Background(), raddr)
}

func (d *tcpDialer) DialContext(ctx context.Context, raddr ma.Multiaddr) (tpt.Conn, error) {
	var c manet.Conn
	var err error
	if d.doReuse {
		c, err = d.reuseDial(ctx, raddr)
	} else {
		c, err = d.madialer.DialContext(ctx, raddr)
	}

	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpDialer) reuseDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	network, netraddr, err := manet.DialArgs(raddr)
	if err != nil {
		return nil, err
	}

	rpev := log.EventBegin(ctx, "tptDialReusePort", logging.LoggableMap{
		"raddr": raddr,
	})

	con, err := d.rd.DialContext(ctx, network, netraddr)
	if err == nil {
		rpev.Done()
		return manet.WrapNetConn(con)
	}
	rpev.SetError(err)
	rpev.Done()

	if !ReuseErrShouldRetry(err) {
		return nil, err
	}

	return d.madialer.DialContext(ctx, raddr)
}

var TCP4 = mafmt.And(mafmt.Base(ma.P_IP4), mafmt.Base(ma.P_TCP))
var TCP6 = mafmt.And(mafmt.Base(ma.P_IP6), mafmt.Base(ma.P_TCP))

func (d *tcpDialer) Matches(a ma.Multiaddr) bool {
	return d.pattern.Matches(a)
}

type tcpListener struct {
	list      manet.Listener
	transport tpt.Transport
}

func (d *tcpListener) Accept() (tpt.Conn, error) {
	c, err := d.list.Accept()
	if err != nil {
		return nil, err
	}

	return &tpt.ConnWrap{
		Conn: c,
		Tpt:  d.transport,
	}, nil
}

func (d *tcpListener) Addr() net.Addr {
	return d.list.Addr()
}

func (t *tcpListener) Multiaddr() ma.Multiaddr {
	return t.list.Multiaddr()
}

func (t *tcpListener) NetListener() net.Listener {
	return t.list.NetListener()
}

func (d *tcpListener) Close() error {
	return d.list.Close()
}
<|endoftext|>


<commit_before>package tcp

import (
	"context"
	"net"
	"runtime"
	"time"

	logging "github.com/ipfs/go-log"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/transport"
	tptu "github.com/libp2p/go-libp2p-transport-upgrader"
	rtpt "github.com/libp2p/go-reuseport-transport"

	ma "github.com/multiformats/go-multiaddr"
	mafmt "github.com/multiformats/go-multiaddr-fmt"
	manet "github.com/multiformats/go-multiaddr/net"
)

// DefaultConnectTimeout is the (default) maximum amount of time the TCP
// transport will spend on the initial TCP connect before giving up.
var DefaultConnectTimeout = 5 * time.Second

var log = logging.Logger("tcp-tpt")

const keepAlivePeriod = 30 * time.Second

type canKeepAlive interface {
	SetKeepAlive(bool) error
	SetKeepAlivePeriod(time.Duration) error
}

var _ canKeepAlive = &net.TCPConn{}

func tryKeepAlive(conn net.Conn, keepAlive bool) {
	keepAliveConn, ok := conn.(canKeepAlive)
	if !ok {
		log.Errorf("Can't set TCP keepalives.")
		return
	}
	if err := keepAliveConn.SetKeepAlive(keepAlive); err != nil {
		log.Errorf("Failed to enable TCP keepalive: %s", err)
		return
	}

	if runtime.GOOS != "openbsd" {
		if err := keepAliveConn.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
			log.Errorf("Failed set keepalive period: %s", err)
		}
	}
}

// try to set linger on the connection, if possible.
func tryLinger(conn net.Conn, sec int) {
	type canLinger interface {
		SetLinger(int) error
	}

	if lingerConn, ok := conn.(canLinger); ok {
		_ = lingerConn.SetLinger(sec)
	}
}

type tcpListener struct {
	manet.Listener
	sec int
}

func (ll *tcpListener) Accept() (manet.Conn, error) {
	c, err := ll.Listener.Accept()
	if err != nil {
		return nil, err
	}
	tryLinger(c, ll.sec)
	tryKeepAlive(c, true)
	return c, nil
}

// TcpTransport is the TCP transport.
type TcpTransport struct {
	// Connection upgrader for upgrading insecure stream connections to
	// secure multiplex connections.
	Upgrader *tptu.Upgrader

	// Explicitly disable reuseport.
	DisableReuseport bool

	// TCP connect timeout
	ConnectTimeout time.Duration

	reuse rtpt.Transport
}

var _ transport.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport(upgrader *tptu.Upgrader) *TcpTransport {
	return &TcpTransport{Upgrader: upgrader, ConnectTimeout: DefaultConnectTimeout}
}

var dialMatcher = mafmt.And(mafmt.IP, mafmt.Base(ma.P_TCP))

// CanDial returns true if this transport believes it can dial the given
// multiaddr.
func (t *TcpTransport) CanDial(addr ma.Multiaddr) bool {
	return dialMatcher.Matches(addr)
}

func (t *TcpTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	// Apply the deadline iff applicable
	if t.ConnectTimeout > 0 {
		deadline := time.Now().Add(t.ConnectTimeout)
		if d, ok := ctx.Deadline(); !ok || deadline.Before(d) {
			var cancel func()
			ctx, cancel = context.WithDeadline(ctx, deadline)
			defer cancel()
		}
	}

	if t.UseReuseport() {
		return t.reuse.DialContext(ctx, raddr)
	}
	var d manet.Dialer
	return d.DialContext(ctx, raddr)
}

// Dial dials the peer at the remote address.
func (t *TcpTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (transport.CapableConn, error) {
	conn, err := t.maDial(ctx, raddr)
	if err != nil {
		return nil, err
	}
	// Set linger to 0 so we never get stuck in the TIME-WAIT state. When
	// linger is 0, connections are _reset_ instead of closed with a FIN.
	// This means we can immediately reuse the 5-tuple and reconnect.
	tryLinger(conn, 0)
	tryKeepAlive(conn, true)
	c, err := newTracingConn(conn, true)
	if err != nil {
		return nil, err
	}
	return t.Upgrader.UpgradeOutbound(ctx, t, c, p)
}

// UseReuseport returns true if reuseport is enabled and available.
func (t *TcpTransport) UseReuseport() bool {
	return !t.DisableReuseport && ReuseportIsAvailable()
}

func (t *TcpTransport) maListen(laddr ma.Multiaddr) (manet.Listener, error) {
	if t.UseReuseport() {
		return t.reuse.Listen(laddr)
	}
	return manet.Listen(laddr)
}

// Listen listens on the given multiaddr.
func (t *TcpTransport) Listen(laddr ma.Multiaddr) (transport.Listener, error) {
	list, err := t.maListen(laddr)
	if err != nil {
		return nil, err
	}
	list = &tracingListener{&tcpListener{list, 0}}
	return t.Upgrader.UpgradeListener(t, list), nil
}

// Protocols returns the list of terminal protocols this transport can dial.
func (t *TcpTransport) Protocols() []int {
	return []int{ma.P_TCP}
}

// Proxy always returns false for the TCP transport.
func (t *TcpTransport) Proxy() bool {
	return false
}

func (t *TcpTransport) String() string {
	return "TCP"
}
<commit_msg>fix: avoid logging "invalid argument" errors when setting keepalive<commit_after>package tcp

import (
	"context"
	"errors"
	"net"
	"os"
	"runtime"
	"time"

	logging "github.com/ipfs/go-log"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/transport"
	tptu "github.com/libp2p/go-libp2p-transport-upgrader"
	rtpt "github.com/libp2p/go-reuseport-transport"

	ma "github.com/multiformats/go-multiaddr"
	mafmt "github.com/multiformats/go-multiaddr-fmt"
	manet "github.com/multiformats/go-multiaddr/net"
)

// DefaultConnectTimeout is the (default) maximum amount of time the TCP
// transport will spend on the initial TCP connect before giving up.
var DefaultConnectTimeout = 5 * time.Second

var log = logging.Logger("tcp-tpt")

const keepAlivePeriod = 30 * time.Second

type canKeepAlive interface {
	SetKeepAlive(bool) error
	SetKeepAlivePeriod(time.Duration) error
}

var _ canKeepAlive = &net.TCPConn{}

func tryKeepAlive(conn net.Conn, keepAlive bool) {
	keepAliveConn, ok := conn.(canKeepAlive)
	if !ok {
		log.Errorf("Can't set TCP keepalives.")
		return
	}
	if err := keepAliveConn.SetKeepAlive(keepAlive); err != nil {
		// Sometimes we seem to get "invalid argument" results from this function on Darwin.
		// This might be due to a closed connection, but I can't reproduce that on Linux.
		//
		// But there's nothing we can do about invalid arguments, so we'll drop this to a
		// debug.
		if errors.Is(err, os.ErrInvalid) {
			log.Debugw("failed to enable TCP keepalive", "error", err)
		} else {
			log.Errorw("failed to enable TCP keepalive", "error", err)
		}
		return
	}

	if runtime.GOOS != "openbsd" {
		if err := keepAliveConn.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
			log.Errorw("failed set keepalive period", "error", err)
		}
	}
}

// try to set linger on the connection, if possible.
func tryLinger(conn net.Conn, sec int) {
	type canLinger interface {
		SetLinger(int) error
	}

	if lingerConn, ok := conn.(canLinger); ok {
		_ = lingerConn.SetLinger(sec)
	}
}

type tcpListener struct {
	manet.Listener
	sec int
}

func (ll *tcpListener) Accept() (manet.Conn, error) {
	c, err := ll.Listener.Accept()
	if err != nil {
		return nil, err
	}
	tryLinger(c, ll.sec)
	tryKeepAlive(c, true)
	return c, nil
}

// TcpTransport is the TCP transport.
type TcpTransport struct {
	// Connection upgrader for upgrading insecure stream connections to
	// secure multiplex connections.
	Upgrader *tptu.Upgrader

	// Explicitly disable reuseport.
	DisableReuseport bool

	// TCP connect timeout
	ConnectTimeout time.Duration

	reuse rtpt.Transport
}

var _ transport.Transport = &TcpTransport{}

// NewTCPTransport creates a tcp transport object that tracks dialers and listeners
// created. It represents an entire tcp stack (though it might not necessarily be)
func NewTCPTransport(upgrader *tptu.Upgrader) *TcpTransport {
	return &TcpTransport{Upgrader: upgrader, ConnectTimeout: DefaultConnectTimeout}
}

var dialMatcher = mafmt.And(mafmt.IP, mafmt.Base(ma.P_TCP))

// CanDial returns true if this transport believes it can dial the given
// multiaddr.
func (t *TcpTransport) CanDial(addr ma.Multiaddr) bool {
	return dialMatcher.Matches(addr)
}

func (t *TcpTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
	// Apply the deadline iff applicable
	if t.ConnectTimeout > 0 {
		deadline := time.Now().Add(t.ConnectTimeout)
		if d, ok := ctx.Deadline(); !ok || deadline.Before(d) {
			var cancel func()
			ctx, cancel = context.WithDeadline(ctx, deadline)
			defer cancel()
		}
	}

	if t.UseReuseport() {
		return t.reuse.DialContext(ctx, raddr)
	}
	var d manet.Dialer
	return d.DialContext(ctx, raddr)
}

// Dial dials the peer at the remote address.
func (t *TcpTransport) Dial(ctx context.Context, raddr ma.Multiaddr, p peer.ID) (transport.CapableConn, error) {
	conn, err := t.maDial(ctx, raddr)
	if err != nil {
		return nil, err
	}
	// Set linger to 0 so we never get stuck in the TIME-WAIT state. When
	// linger is 0, connections are _reset_ instead of closed with a FIN.
	// This means we can immediately reuse the 5-tuple and reconnect.
	tryLinger(conn, 0)
	tryKeepAlive(conn, true)
	c, err := newTracingConn(conn, true)
	if err != nil {
		return nil, err
	}
	return t.Upgrader.UpgradeOutbound(ctx, t, c, p)
}

// UseReuseport returns true if reuseport is enabled and available.
func (t *TcpTransport) UseReuseport() bool {
	return !t.DisableReuseport && ReuseportIsAvailable()
}

func (t *TcpTransport) maListen(laddr ma.Multiaddr) (manet.Listener, error) {
	if t.UseReuseport() {
		return t.reuse.Listen(laddr)
	}
	return manet.Listen(laddr)
}

// Listen listens on the given multiaddr.
func (t *TcpTransport) Listen(laddr ma.Multiaddr) (transport.Listener, error) {
	list, err := t.maListen(laddr)
	if err != nil {
		return nil, err
	}
	list = &tracingListener{&tcpListener{list, 0}}
	return t.Upgrader.UpgradeListener(t, list), nil
}

// Protocols returns the list of terminal protocols this transport can dial.
func (t *TcpTransport) Protocols() []int {
	return []int{ma.P_TCP}
}

// Proxy always returns false for the TCP transport.
func (t *TcpTransport) Proxy() bool {
	return false
}

func (t *TcpTransport) String() string {
	return "TCP"
}
<|endoftext|>


<commit_before>package capnp

const Package = uint64(0xbea97f1023792be0)
const Import = uint64(0xe130b601260e44b5)
const Doc = uint64(0xc58ad6bd519f935e)
const Tag = uint64(0xa574b41924caefc7)
const Notag = uint64(0xc8768679ec52e012)
const Customtype = uint64(0xfa10659ae02f2093)
const Name = uint64(0xc2b96012172f8df1)
<commit_msg>capnp: add doc comment for annotations<commit_after>package capnp

// go.capnp annotation values.
const (
	Package    = uint64(0xbea97f1023792be0)
	Import     = uint64(0xe130b601260e44b5)
	Doc        = uint64(0xc58ad6bd519f935e)
	Tag        = uint64(0xa574b41924caefc7)
	Notag      = uint64(0xc8768679ec52e012)
	Customtype = uint64(0xfa10659ae02f2093)
	Name       = uint64(0xc2b96012172f8df1)
)
<|endoftext|>


<commit_before>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"fmt"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/go/buildbucket"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE_PATH = "infra/branch-config/cq.cfg"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 35
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBots, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, cqTryBots, cqTryBotsFunc, metricName}, err
}

// GetCQTryBots is an interface for returing the CQ trybots of a project.
type GetCQTryBots func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     []string
	cqTryBotsFunc GetCQTryBots
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBots.
func GetSkiaCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBots.
func GetSkiaInfraCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// getCQTryBots is a convenience method for the Skia and Skia Infra CQ TryBots.
func getCQTryBots(repo string) ([]string, error) {
	var buf bytes.Buffer
	if err := gitiles.NewRepo(repo).ReadFile(CQ_CFG_FILE_PATH, &buf); err != nil {
		return nil, err
	}
	var cqCfg Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	tryJobs := []string{}
	for _, bucket := range cqCfg.Verifiers.GetTryJob().GetBuckets() {
		for _, builder := range bucket.GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			tryJobs = append(tryJobs, builder.GetName())
		}
	}
	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}

	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	c.cqTryBots = tryBots
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
func (c *Client) ReportCQStats(change int64) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucket.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Parameters.BuilderName) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime == time.Time{}.UTC()) {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Parameters.BuilderName, gerritURL)
			continue
		}
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Parameters.BuilderName, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION), durationTags)
		landedTrybotDurationMetric.Update(duration)

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime != time.Time{}.UTC()) {
			if time.Hour*24 < time.Now().UTC().Sub(createdTime) {
				// The build has completed more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Parameters.BuilderName, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION), durationTags)
		inflightTrybotDurationMetric.Update(duration)

	}
	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := util.ContainsAny(builderName, c.cqTryBots)
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<commit_msg>[CQ Watcher] Use the right membership function in isCQTryBot<commit_after>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"fmt"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"go.skia.org/infra/go/sklog"

	"go.skia.org/infra/go/buildbucket"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE_PATH = "infra/branch-config/cq.cfg"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 35
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBots, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, cqTryBots, cqTryBotsFunc, metricName}, err
}

// GetCQTryBots is an interface for returing the CQ trybots of a project.
type GetCQTryBots func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     []string
	cqTryBotsFunc GetCQTryBots
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBots.
func GetSkiaCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBots.
func GetSkiaInfraCQTryBots() ([]string, error) {
	return getCQTryBots(common.REPO_SKIA)
}

// getCQTryBots is a convenience method for the Skia and Skia Infra CQ TryBots.
func getCQTryBots(repo string) ([]string, error) {
	var buf bytes.Buffer
	if err := gitiles.NewRepo(repo).ReadFile(CQ_CFG_FILE_PATH, &buf); err != nil {
		return nil, err
	}
	var cqCfg Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	tryJobs := []string{}
	for _, bucket := range cqCfg.Verifiers.GetTryJob().GetBuckets() {
		for _, builder := range bucket.GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			tryJobs = append(tryJobs, builder.GetName())
		}
	}
	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}

	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	c.cqTryBots = tryBots
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
func (c *Client) ReportCQStats(change int64) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucket.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Parameters.BuilderName) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime == time.Time{}.UTC()) {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Parameters.BuilderName, gerritURL)
			continue
		}
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Parameters.BuilderName, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION), durationTags)
		landedTrybotDurationMetric.Update(duration)

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucket.Build, gerritURL string) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime := time.Time(b.Created).UTC()
		completedTime := time.Time(b.Completed).UTC()
		if (completedTime != time.Time{}.UTC()) {
			if time.Hour*24 < time.Now().UTC().Sub(createdTime) {
				// The build has completed more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		durationTags := map[string]string{
			"trybot": b.Parameters.BuilderName,
		}
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Parameters.BuilderName, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION), durationTags)
		inflightTrybotDurationMetric.Update(duration)

	}
	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := util.In(builderName, c.cqTryBots)
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<|endoftext|>


<commit_before>package main

//author: Doug Watson

//macOS will need more open files than the default 256 if you want to run over a couple hundred goroutines
//launchctl limit maxfiles 10200

//For a really big test, if you need a million open files on a mac:
//nvram boot-args="serverperfmode=1"
//shutdown -r now
//launchctl limit maxfiles 999999
//ulimit -n 999999

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"
	"syscall"
	"time"

	runner "github.com/adt-automation/goRunner/golib"
)

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

var (
	clients        int
	secsTimedTest  int
	noRampUp       bool
	debug          bool
	dump           bool
	connectTimeout int
	writeTimeout   int
	readTimeout    int
	targetTPS      float64
	baseUrl        string
	configFile     string
	inputFile      string
	inputFile1     string
	delimeter      string
	headerExit     bool
	noHeader       bool

//	wg             sync.WaitGroup
)

type MyConn struct {
	net.Conn
	readTimeout  time.Duration
	writeTimeout time.Duration
	result       *runner.Result
}

func init() {
	flag.IntVar(&clients, "c", 100, "Number of concurrent clients to launch.")
	flag.IntVar(&secsTimedTest, "t", 0, "Time in seconds for a timed load test")
	flag.StringVar(&inputFile, "f", "", "Read input from file rather than stdin")
	flag.StringVar(&inputFile1, "fh", "", "Read input from csv file which has a header row")
	flag.BoolVar(&noRampUp, "nr", false, "Do not ramp up to total concurrent number of clients")
	flag.Float64Var(&targetTPS, "targetTPS", 1000000, "The default max TPS is set to 1 million. Good luck reaching this :p")
	flag.StringVar(&baseUrl, "baseUrl", "", "The host to test. Example https://test2.someserver.org")
	flag.BoolVar(&debug, "debug", false, "Show the body returned for the api call")
	flag.BoolVar(&dump, "dump", false, "Show the HTTP dump for the api call")
	flag.StringVar(&configFile, "configFile", "config.ini", "Config file location")
	flag.StringVar(&delimeter, "d", ",", "Output file delimeter")
	flag.BoolVar(&headerExit, "hx", false, "Print output header row and exit")
	flag.BoolVar(&noHeader, "nh", false, "Suppress output header row. Ignored if hx is set")
	flag.IntVar(&readTimeout, "readtimeout", 30, "Timeout in seconds for the target API to send the first response byte. Default 30 seconds")
}

/*
func generateInput() chan string {
	lines := make(chan string)
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			lines <- scanner.Text()
		}
		//scanner.Err()
		fmt.Printf("scanner done:\n\n")
	}()
	return lines
}
*/
func main() {
	flag.Parse()
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Println(err)
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	runner.Delimeter = delimeter //set the delimeter we want for the output log in the imported runner package
	if !headerExit && baseUrl == "" {
		fmt.Fprintf(os.Stderr, "\nPlease provide the baseUrl...\n\n")
		flag.Usage()
		os.Exit(1)
	}
	trafficChannel := make(chan string)
	//	startTraffic(trafficChannel) //start reading on the channel

	if *runner.Verbose {
		println("verbose")
	}

	overallStartTime := time.Now()
	baseUrlFilter := regexp.MustCompile(baseUrl)
	configuration := runner.NewConfiguration2(configFile)
	results := make(map[int]*runner.Result)
	signalChannel := make(chan os.Signal, 2)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
	go func() {
		_ = <-signalChannel
		runner.ExitWithStatus(results, overallStartTime)
	}()
	goMaxProcs := os.Getenv("GOMAXPROCS")
	if goMaxProcs == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	if !headerExit {
		printSessionSummary(configuration, configFile)
	}

	var rampUpDelayMs time.Duration = time.Duration(0)
	if !headerExit && !noRampUp && clients > 1 {
		rampUpDelayMs = calcRampUpDelay(configuration)
		fmt.Fprintf(os.Stderr, "Spacing sessions %v apart to ramp up to %d client sessions\n", rampUpDelayMs, clients)
	}
	transportConfig := &tls.Config{InsecureSkipVerify: true} //allow wrong ssl certs
	tr := &http.Transport{TLSClientConfig: transportConfig}
	tr.ResponseHeaderTimeout = time.Second * time.Duration(readTimeout)
	var done sync.WaitGroup
	var stopTime time.Time // client loops will work to the end of trafficChannel unless we explicitly init a stopTime
	if secsTimedTest > 0 {
		stopTime = time.Now().Add(time.Duration(secsTimedTest) * time.Second)
	}
	for i := 0; i < clients; i++ {
		done.Add(1)
		result := &runner.Result{}
		results[i] = result
		clientDelay := rampUpDelayMs * time.Duration(i)
		go client(tr, configuration, result, &done, trafficChannel, i, baseUrlFilter, stopTime, clientDelay)
	}

	if len(inputFile1) > 0 {
		if len(inputFile) > 0 {
			fmt.Fprintf(os.Stderr, "don't use -f and -fh together\n")
			os.Exit(1)
		} else {
			inputFile = inputFile1
		}
	}

	scanner := bufio.NewScanner(os.Stdin)
	if len(inputFile) > 0 {
		file, err := os.Open(inputFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err.Error())
			os.Exit(1)
		}
		defer file.Close()
		scanner = bufio.NewScanner(file)
	}

	_ = scanner.Scan()
	if !noHeader {
		runner.PrintLogHeader(scanner.Text(), len(inputFile1) > 0)
		if headerExit {
			os.Exit(0)
		}
	}
	if len(inputFile1) > 0 && runner.HasInputColHeaders() == false {
		// name our column headers
		runner.HeadInputColumns(scanner.Text())
	} else {
		// anonymous column headers, put work into the channel instead
		trafficChannel <- scanner.Text()
	}

	for scanner.Scan() {
		trafficChannel <- scanner.Text() //put work into the channel from Stdin
	}
	close(trafficChannel)
	done.Wait()
	runner.ExitWithStatus(results, overallStartTime)
}

func calcRampUpDelay(cfg *runner.CfgStruct) time.Duration {
	if clients == 0 {
		return time.Duration(0)
	} else {
		dur := runner.EstimateSessionTime(cfg)
		return dur / time.Duration(clients)
	}
}

//one client per -c=thread
//func noRedirect(req *http.Request, via []*http.Request) error {
//	return errors.New("Don't redirect!")
//}
func client(tr *http.Transport, configuration *runner.CfgStruct, result *runner.Result, done *sync.WaitGroup, trafficChannel chan string, clientId int, baseUrlFilter *regexp.Regexp, stopTime time.Time, initialDelay time.Duration) {
	defer done.Done()
	// strangely, this sleep cannot be moved outside the client function
	// or all client threads will wait until the last one is spawned before starting their own execution loops
	time.Sleep(initialDelay)

	msDelay := 0
	grep1, grep2 := "", ""
	for id := range trafficChannel {
		//read in the line. either a generated id or (TODO) a url
		//		id := <-trafficChannel //urlTemplate : 100|10000|about

		//Do Heartbeat
		// cookieMap and sessionVars should start fresh every time we start a DoReq session
		var cookieMap = make(map[string]*http.Cookie)
		var sessionVars = make(map[string]string)
		runner.DoReq(0, id, configuration, result, clientId, baseUrl, baseUrlFilter, msDelay, tr, cookieMap, sessionVars, grep1, grep2, stopTime, 0.0) //val,resp, err
		msDelay = runner.PostSessionDelay
	}
	if time.Now().Before(stopTime) {
		fmt.Fprintf(os.Stderr, "client %d ran out of test input %.2fs before full test time\n", clientId, stopTime.Sub(time.Now()).Seconds())
	}
}

func printSessionSummary(configuration *runner.CfgStruct, configFile string) {
	fmt.Fprintf(os.Stderr, "GORUNNER\n")
	if len(configuration.Version.ConfigVersion) > 5 {
		fmt.Fprintf(os.Stderr, "Configuration:                  %s version %s\n", configFile, configuration.Version.ConfigVersion[5:len(configuration.Version.ConfigVersion)-1])
	} else {
		fmt.Fprintf(os.Stderr, "Configuration:                  %s\n", configFile)
	}
	fmt.Fprintf(os.Stderr, "Session profile:                %d requests: %s\n", len(runner.CommandQueue), strings.Join(runner.CommandQueue, ", "))
	fmt.Fprintf(os.Stderr, "Estimated session time:         %v\n", runner.EstimateSessionTime(configuration))
	fmt.Fprintf(os.Stderr, "Simultaneous sessions:          %d\n", clients)
	fmt.Fprintf(os.Stderr, "API host:                       %s\n", baseUrl)
	const layout = "2006-01-02 15:04:05"
	fmt.Fprintf(os.Stderr, "Test start time:                %v\n", time.Now().Format(layout))
}
<commit_msg>added rampup flag, not tested<commit_after>package main

//author: Doug Watson

//macOS will need more open files than the default 256 if you want to run over a couple hundred goroutines
//launchctl limit maxfiles 10200

//For a really big test, if you need a million open files on a mac:
//nvram boot-args="serverperfmode=1"
//shutdown -r now
//launchctl limit maxfiles 999999
//ulimit -n 999999

import (
	"bufio"
	"crypto/tls"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"
	"syscall"
	"time"

	runner "github.com/adt-automation/goRunner/golib"
)

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

var (
	clients        int
	secsTimedTest  int
	debug          bool
	dump           bool
	connectTimeout int
	writeTimeout   int
	readTimeout    int
	targetTPS      float64
	baseUrl        string
	configFile     string
	inputFile      string
	inputFile1     string
	delimeter      string
	headerExit     bool
	noHeader       bool
	rampUp         string

//	wg             sync.WaitGroup
)

type MyConn struct {
	net.Conn
	readTimeout  time.Duration
	writeTimeout time.Duration
	result       *runner.Result
}

func init() {
	flag.IntVar(&clients, "c", 100, "Number of concurrent clients to launch.")
	flag.IntVar(&secsTimedTest, "t", 0, "Time in seconds for a timed load test")
	flag.StringVar(&inputFile, "f", "", "Read input from file rather than stdin")
	flag.StringVar(&inputFile1, "fh", "", "Read input from csv file which has a header row")
	// flag.BoolVar(&noRampUp, "nr", false, "Do not ramp up to total concurrent number of clients")
	flag.StringVar(&rampUp, "rampUp", "0", "Specify ramp up delay as duration (1m2s, 300ms, ..). 'auto' will compute from client sessions. Default to 0 (no ramp up)")
	flag.Float64Var(&targetTPS, "targetTPS", 1000000, "The default max TPS is set to 1 million. Good luck reaching this :p")
	flag.StringVar(&baseUrl, "baseUrl", "", "The host to test. Example https://test2.someserver.org")
	flag.BoolVar(&debug, "debug", false, "Show the body returned for the api call")
	flag.BoolVar(&dump, "dump", false, "Show the HTTP dump for the api call")
	flag.StringVar(&configFile, "configFile", "config.ini", "Config file location")
	flag.StringVar(&delimeter, "d", ",", "Output file delimeter")
	flag.BoolVar(&headerExit, "hx", false, "Print output header row and exit")
	flag.BoolVar(&noHeader, "nh", false, "Suppress output header row. Ignored if hx is set")
	flag.IntVar(&readTimeout, "readtimeout", 30, "Timeout in seconds for the target API to send the first response byte. Default 30 seconds")

}

/*
func generateInput() chan string {
	lines := make(chan string)
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			lines <- scanner.Text()
		}
		//scanner.Err()
		fmt.Printf("scanner done:\n\n")
	}()
	return lines
}
*/
func main() {
	// ---------------------------------------------------------------------------------------------
	// Parse flags
	flag.Parse()

	// ---------------------------------------------------------------------------------------------
	// Validate input
	if clients < 1 {
		fmt.Fprintf(os.Stderr, "\nNumber of concurrent client should be at least 1\n\n")
		flag.Usage()
		os.Exit(1)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			fmt.Println(err)
			return
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	runner.Delimeter = delimeter //set the delimeter we want for the output log in the imported runner package
	if !headerExit && baseUrl == "" {
		fmt.Fprintf(os.Stderr, "\nPlease provide the baseUrl...\n\n")
		flag.Usage()
		os.Exit(1)
	}
	trafficChannel := make(chan string)
	//	startTraffic(trafficChannel) //start reading on the channel

	if *runner.Verbose {
		println("verbose")
	}

	overallStartTime := time.Now()
	baseUrlFilter := regexp.MustCompile(baseUrl)
	configuration := runner.NewConfiguration2(configFile)
	results := make(map[int]*runner.Result)
	signalChannel := make(chan os.Signal, 2)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
	go func() {
		_ = <-signalChannel
		runner.ExitWithStatus(results, overallStartTime)
	}()
	goMaxProcs := os.Getenv("GOMAXPROCS")
	if goMaxProcs == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	rampUpDelay := time.Duration(0)
	if !headerExit {
		printSessionSummary(configuration, configFile)
		if rampUp == "auto" {
			rampUpDelay = calcRampUpDelay(configuration)
		} else {
			var err error
			rampUpDelay, err = time.ParseDuration(rampUp)
			if err != nil {
				fmt.Println("invalid rampUp duration" + err.Error())
				os.Exit(1)
			}
		}
	}
	fmt.Fprintf(os.Stderr, "Spacing sessions %v apart to ramp up to %d client sessions\n", rampUpDelay, clients)

	transportConfig := &tls.Config{InsecureSkipVerify: true} //allow wrong ssl certs
	tr := &http.Transport{TLSClientConfig: transportConfig}
	tr.ResponseHeaderTimeout = time.Second * time.Duration(readTimeout)
	var done sync.WaitGroup
	var stopTime time.Time // client loops will work to the end of trafficChannel unless we explicitly init a stopTime
	if secsTimedTest > 0 {
		stopTime = time.Now().Add(time.Duration(secsTimedTest) * time.Second)
	}
	for i := 0; i < clients; i++ {
		done.Add(1)
		result := &runner.Result{}
		results[i] = result
		clientDelay := rampUpDelay * time.Duration(i)
		go client(tr, configuration, result, &done, trafficChannel, i, baseUrlFilter, stopTime, clientDelay)
	}

	if len(inputFile1) > 0 {
		if len(inputFile) > 0 {
			fmt.Fprintf(os.Stderr, "don't use -f and -fh together\n")
			os.Exit(1)
		} else {
			inputFile = inputFile1
		}
	}

	scanner := bufio.NewScanner(os.Stdin)
	if len(inputFile) > 0 {
		file, err := os.Open(inputFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err.Error())
			os.Exit(1)
		}
		defer file.Close()
		scanner = bufio.NewScanner(file)
	}

	_ = scanner.Scan()
	if !noHeader {
		runner.PrintLogHeader(scanner.Text(), len(inputFile1) > 0)
		if headerExit {
			os.Exit(0)
		}
	}
	if len(inputFile1) > 0 && runner.HasInputColHeaders() == false {
		// name our column headers
		runner.HeadInputColumns(scanner.Text())
	} else {
		// anonymous column headers, put work into the channel instead
		trafficChannel <- scanner.Text()
	}

	for scanner.Scan() {
		trafficChannel <- scanner.Text() //put work into the channel from Stdin
	}
	close(trafficChannel)
	done.Wait()
	runner.ExitWithStatus(results, overallStartTime)
}

func calcRampUpDelay(cfg *runner.CfgStruct) time.Duration {
	if clients == 0 {
		return time.Duration(0)
	} else {
		dur := runner.EstimateSessionTime(cfg)
		return dur / time.Duration(clients)
	}
}

//one client per -c=thread
//func noRedirect(req *http.Request, via []*http.Request) error {
//	return errors.New("Don't redirect!")
//}
func client(tr *http.Transport, configuration *runner.CfgStruct, result *runner.Result, done *sync.WaitGroup, trafficChannel chan string, clientId int, baseUrlFilter *regexp.Regexp, stopTime time.Time, initialDelay time.Duration) {
	defer done.Done()
	// strangely, this sleep cannot be moved outside the client function
	// or all client threads will wait until the last one is spawned before starting their own execution loops
	time.Sleep(initialDelay)

	msDelay := 0
	grep1, grep2 := "", ""
	for id := range trafficChannel {
		//read in the line. either a generated id or (TODO) a url
		//		id := <-trafficChannel //urlTemplate : 100|10000|about

		//Do Heartbeat
		// cookieMap and sessionVars should start fresh every time we start a DoReq session
		var cookieMap = make(map[string]*http.Cookie)
		var sessionVars = make(map[string]string)
		runner.DoReq(0, id, configuration, result, clientId, baseUrl, baseUrlFilter, msDelay, tr, cookieMap, sessionVars, grep1, grep2, stopTime, 0.0) //val,resp, err
		msDelay = runner.PostSessionDelay
	}
	if time.Now().Before(stopTime) {
		fmt.Fprintf(os.Stderr, "client %d ran out of test input %.2fs before full test time\n", clientId, stopTime.Sub(time.Now()).Seconds())
	}
}

func printSessionSummary(configuration *runner.CfgStruct, configFile string) {
	fmt.Fprintf(os.Stderr, "GORUNNER\n")
	if len(configuration.Version.ConfigVersion) > 5 {
		fmt.Fprintf(os.Stderr, "Configuration:                  %s version %s\n", configFile, configuration.Version.ConfigVersion[5:len(configuration.Version.ConfigVersion)-1])
	} else {
		fmt.Fprintf(os.Stderr, "Configuration:                  %s\n", configFile)
	}
	fmt.Fprintf(os.Stderr, "Session profile:                %d requests: %s\n", len(runner.CommandQueue), strings.Join(runner.CommandQueue, ", "))
	fmt.Fprintf(os.Stderr, "Estimated session time:         %v\n", runner.EstimateSessionTime(configuration))
	fmt.Fprintf(os.Stderr, "Simultaneous sessions:          %d\n", clients)
	fmt.Fprintf(os.Stderr, "API host:                       %s\n", baseUrl)
	const layout = "2006-01-02 15:04:05"
	fmt.Fprintf(os.Stderr, "Test start time:                %v\n", time.Now().Format(layout))
}
<|endoftext|>