// Copyright (C) 2019 Storj Labs, Inc. // See LICENSE for copying information. // See LICENSE for copying information. package kademlia import ( "bytes" "context" "io/ioutil" "math/rand" "net" "os" "path/filepath" "strconv" "sync/atomic" "testing" "time" "github.com/stretchr/testify/assert" "go.uber.org/zap/zaptest" "google.golang.org/grpc" "storj.io/storj/internal/testcontext" "storj.io/storj/internal/testidentity" "storj.io/storj/internal/teststorj" "storj.io/storj/pkg/identity" "storj.io/storj/pkg/node" "storj.io/storj/pkg/pb" "storj.io/storj/pkg/provider" "storj.io/storj/pkg/storj" ) const ( defaultAlpha = 5 ) func TestNewKademlia(t *testing.T) { ctx := testcontext.New(t) defer ctx.Cleanup() rootdir, cleanup := mktempdir(t, "kademlia") defer cleanup() cases := []struct { id *identity.FullIdentity bn []pb.Node addr string expectedErr error }{ { id: func() *identity.FullIdentity { id, err := testidentity.NewTestIdentity(ctx) assert.NoError(t, err) return id }(), bn: []pb.Node{{Id: teststorj.NodeIDFromString("foo")}}, addr: "127.0.0.1:8080", }, { id: func() *provider.FullIdentity { id, err := testidentity.NewTestIdentity(ctx) assert.NoError(t, err) return id }(), bn: []pb.Node{{Id: teststorj.NodeIDFromString("foo")}}, addr: "127.0.0.1:8080", }, } for i, v := range cases { dir := filepath.Join(rootdir, strconv.Itoa(i)) kad, err := NewKademlia(zaptest.NewLogger(t), pb.NodeType_STORAGE, v.bn, v.addr, nil, v.id, dir, defaultAlpha) assert.NoError(t, err) assert.Equal(t, v.expectedErr, err) assert.Equal(t, kad.bootstrapNodes, v.bn) assert.NotNil(t, kad.dialer) assert.NotNil(t, kad.routingTable) assert.NoError(t, kad.Disconnect()) } } func TestPeerDiscovery(t *testing.T) { ctx := testcontext.New(t) defer ctx.Cleanup() dir, cleanup := mktempdir(t, "kademlia") defer cleanup() // make new identity bootServer, mockBootServer, bootID, bootAddress := startTestNodeServer(ctx) defer bootServer.Stop() testServer, _, testID, testAddress := startTestNodeServer(ctx) defer testServer.Stop() targetServer, _, targetID, targetAddress := startTestNodeServer(ctx) defer targetServer.Stop() bootstrapNodes := []pb.Node{{Id: bootID.ID, Address: &pb.NodeAddress{Address: bootAddress}, Type: pb.NodeType_STORAGE}} metadata := &pb.NodeMetadata{ Email: "foo@bar.com", Wallet: "OperatorWallet", } k, err := NewKademlia(zaptest.NewLogger(t), pb.NodeType_STORAGE, bootstrapNodes, testAddress, metadata, testID, dir, defaultAlpha) assert.NoError(t, err) rt, err := k.GetRoutingTable(ctx) assert.NoError(t, err) assert.Equal(t, rt.Local().Metadata.Email, "foo@bar.com") assert.Equal(t, rt.Local().Metadata.Wallet, "OperatorWallet") defer func() { assert.NoError(t, k.Disconnect()) }() cases := []struct { target storj.NodeID expected *pb.Node expectedErr error }{ {target: func() storj.NodeID { mockBootServer.returnValue = []*pb.Node{{Id: targetID.ID, Type: pb.NodeType_STORAGE, Address: &pb.NodeAddress{Address: targetAddress}}} return targetID.ID }(), expected: &pb.Node{}, expectedErr: nil, }, {target: bootID.ID, expected: nil, expectedErr: nil, }, } for _, v := range cases { _, err := k.lookup(ctx, v.target, true) assert.Equal(t, v.expectedErr, err) } } func TestBootstrap(t *testing.T) { ctx := testcontext.New(t) defer ctx.Cleanup() bn, s, clean := testNode(t, []pb.Node{}) defer clean() defer s.Stop() n1, s1, clean1 := testNode(t, []pb.Node{bn.routingTable.self}) defer clean1() defer s1.Stop() err := n1.Bootstrap(ctx) assert.NoError(t, err) n2, s2, clean2 := testNode(t, []pb.Node{bn.routingTable.self}) defer clean2() defer s2.Stop() err = n2.Bootstrap(ctx) assert.NoError(t, err) nodeIDs, err := n2.routingTable.nodeBucketDB.List(nil, 0) assert.NoError(t, err) assert.Len(t, nodeIDs, 3) } func testNode(t *testing.T, bn []pb.Node) (*Kademlia, *grpc.Server, func()) { ctx := testcontext.New(t) // new address lis, err := net.Listen("tcp", "127.0.0.1:0") assert.NoError(t, err) // new config // new identity fid, err := testidentity.NewTestIdentity(ctx) assert.NoError(t, err) // new kademlia dir, cleanup := mktempdir(t, "kademlia") logger := zaptest.NewLogger(t) k, err := NewKademlia(logger, pb.NodeType_STORAGE, bn, lis.Addr().String(), nil, fid, dir, defaultAlpha) assert.NoError(t, err) s := node.NewServer(logger, k) // new ident opts identOpt, err := fid.ServerOption() assert.NoError(t, err) grpcServer := grpc.NewServer(identOpt) pb.RegisterNodesServer(grpcServer, s) go func() { assert.NoError(t, grpcServer.Serve(lis)) }() return k, grpcServer, func() { defer cleanup() assert.NoError(t, k.Disconnect()) } } func TestRefresh(t *testing.T) { ctx := testcontext.New(t) defer ctx.Cleanup() k, s, clean := testNode(t, []pb.Node{}) defer clean() defer s.Stop() //turn back time for only bucket rt := k.routingTable now := time.Now().UTC() bID := firstBucketID //always exists err := rt.SetBucketTimestamp(bID[:], now.Add(-2*time.Hour)) assert.NoError(t, err) //refresh should call FindNode, updating the time err = k.refresh(ctx) assert.NoError(t, err) ts1, err := rt.GetBucketTimestamp(bID[:]) assert.NoError(t, err) assert.True(t, now.Add(-5*time.Minute).Before(ts1)) //refresh should not call FindNode, leaving the previous time err = k.refresh(ctx) assert.NoError(t, err) ts2, err := rt.GetBucketTimestamp(bID[:]) assert.NoError(t, err) assert.True(t, ts1.Equal(ts2)) } func TestGetNodes(t *testing.T) { ctx := testcontext.New(t) defer ctx.Cleanup() var ( nodeIDA = teststorj.NodeIDFromString("AAAAA") nodeIDB = teststorj.NodeIDFromString("BBBBB") nodeIDC = teststorj.NodeIDFromString("CCCCC") nodeIDD = teststorj.NodeIDFromString("DDDDD") ) lis, err := net.Listen("tcp", "127.0.0.1:0") assert.NoError(t, err) srv, _ := newTestServer(ctx, []*pb.Node{{Id: teststorj.NodeIDFromString("foo")}}) go func() { assert.NoError(t, srv.Serve(lis)) }() defer srv.Stop() // make new identity fid, err := testidentity.NewTestIdentity(ctx) assert.NoError(t, err) fid2, err := testidentity.NewTestIdentity(ctx) assert.NoError(t, err) fid.ID = nodeIDA fid2.ID = nodeIDB // create two new unique identities assert.NotEqual(t, fid.ID, fid2.ID) dir, cleanup := mktempdir(t, "kademlia") defer cleanup() k, err := NewKademlia(zaptest.NewLogger(t), pb.NodeType_STORAGE, []pb.Node{{Id: fid2.ID, Address: &pb.NodeAddress{Address: lis.Addr().String()}}}, lis.Addr().String(), nil, fid, dir, defaultAlpha) assert.NoError(t, err) defer func() { assert.NoError(t, k.Disconnect()) }() // add nodes ids := storj.NodeIDList{nodeIDA, nodeIDB, nodeIDC, nodeIDD} bw := []int64{1, 2, 3, 4} disk := []int64{4, 3, 2, 1} nodes := []*pb.Node{} for i, v := range ids { n := &pb.Node{ Id: v, Restrictions: &pb.NodeRestrictions{ FreeBandwidth: bw[i], FreeDisk: disk[i], }, Type: pb.NodeType_STORAGE, } nodes = append(nodes, n) err = k.routingTable.ConnectionSuccess(n) assert.NoError(t, err) } cases := []struct { testID string start storj.NodeID limit int restrictions []pb.Restriction expected []*pb.Node }{ {testID: "one", start: nodeIDB, limit: 2, restrictions: []pb.Restriction{ { Operator: pb.Restriction_GT, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(2), }, }, expected: nodes[2:], }, {testID: "two", start: nodeIDA, limit: 3, restrictions: []pb.Restriction{ { Operator: pb.Restriction_GT, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(2), }, { Operator: pb.Restriction_LT, Operand: pb.Restriction_FREE_DISK, Value: int64(2), }, }, expected: nodes[3:], }, {testID: "three", start: nodeIDA, limit: 4, restrictions: []pb.Restriction{}, expected: nodes, }, } for _, c := range cases { t.Run(c.testID, func(t *testing.T) { ns, err := k.GetNodes(ctx, c.start, c.limit, c.restrictions...) assert.NoError(t, err) assert.Equal(t, len(c.expected), len(ns)) for i, n := range ns { assert.True(t, bytes.Equal(c.expected[i].Id.Bytes(), n.Id.Bytes())) } }) } } func TestMeetsRestrictions(t *testing.T) { cases := []struct { testID string r []pb.Restriction n pb.Node expect bool }{ {testID: "pass one", r: []pb.Restriction{ { Operator: pb.Restriction_EQ, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(1), }, }, n: pb.Node{ Restrictions: &pb.NodeRestrictions{ FreeBandwidth: int64(1), }, }, expect: true, }, {testID: "pass multiple", r: []pb.Restriction{ { Operator: pb.Restriction_LTE, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(2), }, { Operator: pb.Restriction_GTE, Operand: pb.Restriction_FREE_DISK, Value: int64(2), }, }, n: pb.Node{ Restrictions: &pb.NodeRestrictions{ FreeBandwidth: int64(1), FreeDisk: int64(3), }, }, expect: true, }, {testID: "fail one", r: []pb.Restriction{ { Operator: pb.Restriction_LT, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(2), }, { Operator: pb.Restriction_GT, Operand: pb.Restriction_FREE_DISK, Value: int64(2), }, }, n: pb.Node{ Restrictions: &pb.NodeRestrictions{ FreeBandwidth: int64(2), FreeDisk: int64(3), }, }, expect: false, }, {testID: "fail multiple", r: []pb.Restriction{ { Operator: pb.Restriction_LT, Operand: pb.Restriction_FREE_BANDWIDTH, Value: int64(2), }, { Operator: pb.Restriction_GT, Operand: pb.Restriction_FREE_DISK, Value: int64(2), }, }, n: pb.Node{ Restrictions: &pb.NodeRestrictions{ FreeBandwidth: int64(2), FreeDisk: int64(2), }, }, expect: false, }, } for _, c := range cases { t.Run(c.testID, func(t *testing.T) { result := meetsRestrictions(c.r, c.n) assert.Equal(t, c.expect, result) }) } } func mktempdir(t *testing.T, dir string) (string, func()) { rootdir, err := ioutil.TempDir("", dir) assert.NoError(t, err) cleanup := func() { assert.NoError(t, os.RemoveAll(rootdir)) } return rootdir, cleanup } func startTestNodeServer(ctx context.Context) (*grpc.Server, *mockNodesServer, *provider.FullIdentity, string) { lis, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { return nil, nil, nil, "" } ca, err := testidentity.NewTestCA(ctx) if err != nil { return nil, nil, nil, "" } identity, err := ca.NewIdentity() if err != nil { return nil, nil, nil, "" } identOpt, err := identity.ServerOption() if err != nil { return nil, nil, nil, "" } grpcServer := grpc.NewServer(identOpt) mn := &mockNodesServer{queryCalled: 0} pb.RegisterNodesServer(grpcServer, mn) go func() { if err := grpcServer.Serve(lis); err != nil { return } }() return grpcServer, mn, identity, lis.Addr().String() } func newTestServer(ctx context.Context, nn []*pb.Node) (*grpc.Server, *mockNodesServer) { ca, err := testidentity.NewTestCA(ctx) if err != nil { return nil, nil } identity, err := ca.NewIdentity() if err != nil { return nil, nil } identOpt, err := identity.ServerOption() if err != nil { return nil, nil } grpcServer := grpc.NewServer(identOpt) mn := &mockNodesServer{queryCalled: 0} pb.RegisterNodesServer(grpcServer, mn) return grpcServer, mn } // TestRandomIds makes sure finds a random node ID is within a range (start..end] func TestRandomIds(t *testing.T) { for x := 0; x < 1000; x++ { var start, end bucketID // many valid options rand.Read(start[:]) rand.Read(end[:]) if bytes.Compare(start[:], end[:]) > 0 { start, end = end, start } id, err := randomIDInRange(start, end) assert.NoError(t, err, "Unexpected err in randomIDInRange") assert.True(t, bytes.Compare(id[:], start[:]) > 0, "Random id was less than starting id") assert.True(t, bytes.Compare(id[:], end[:]) <= 0, "Random id was greater than end id") //invalid range _, err = randomIDInRange(end, start) assert.Error(t, err, "Missing expected err in invalid randomIDInRange") //no valid options end = start _, err = randomIDInRange(start, end) assert.Error(t, err, "Missing expected err in empty randomIDInRange") // one valid option if start[31] == 255 { start[31] = 254 } else { end[31] = start[31] + 1 } id, err = randomIDInRange(start, end) assert.NoError(t, err, "Unexpected err in randomIDInRange") assert.True(t, bytes.Equal(id[:], end[:]), "Not-so-random id was incorrect") } } type mockNodesServer struct { queryCalled int32 pingCalled int32 returnValue []*pb.Node } func (mn *mockNodesServer) Query(ctx context.Context, req *pb.QueryRequest) (*pb.QueryResponse, error) { atomic.AddInt32(&mn.queryCalled, 1) return &pb.QueryResponse{Response: mn.returnValue}, nil } func (mn *mockNodesServer) Ping(ctx context.Context, req *pb.PingRequest) (*pb.PingResponse, error) { atomic.AddInt32(&mn.pingCalled, 1) return &pb.PingResponse{}, nil }