move bitmap into separate package

This commit is contained in:
Alex Grintsvayg 2018-06-14 11:48:02 -04:00
parent 5968953d05
commit fc9b05b8c6
16 changed files with 243 additions and 220 deletions

View file

@ -1,38 +1,38 @@
package dht package bits
import ( import (
"bytes"
"crypto/rand" "crypto/rand"
"encoding/hex" "encoding/hex"
"strconv"
"strings" "strings"
"strconv"
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
) )
// TODO: http://roaringbitmap.org/ // TODO: http://roaringbitmap.org/
const (
NumBytes = 48 // bytes
NumBits = NumBytes * 8
)
// Bitmap is a generalized representation of an identifier or data that can be sorted, compared fast. Used by the DHT // Bitmap is a generalized representation of an identifier or data that can be sorted, compared fast. Used by the DHT
// package as a way to handle the unique identifiers of a DHT node. // package as a way to handle the unique identifiers of a DHT node.
type Bitmap [nodeIDLength]byte type Bitmap [NumBytes]byte
func (b Bitmap) rawString() string { func (b Bitmap) String() string {
return string(b[:]) return string(b[:])
} }
// BString returns the bitmap as a string of 0s and 1s // BString returns the bitmap as a string of 0s and 1s
func (b Bitmap) BString() string { func (b Bitmap) BString() string {
var buf bytes.Buffer var s string
for i := 0; i < nodeIDBits; i++ { for _, byte := range b {
if b.Get(i) { s += strconv.FormatInt(int64(byte), 2)
buf.WriteString("1")
} else {
buf.WriteString("0")
} }
} return s
return buf.String()
} }
// Hex returns a hexadecimal representation of the bitmap. // Hex returns a hexadecimal representation of the bitmap.
@ -147,7 +147,7 @@ func (b Bitmap) Not() Bitmap {
func (b Bitmap) add(other Bitmap) (Bitmap, bool) { func (b Bitmap) add(other Bitmap) (Bitmap, bool) {
var ret Bitmap var ret Bitmap
carry := false carry := false
for i := nodeIDBits - 1; i >= 0; i-- { for i := NumBits - 1; i >= 0; i-- {
bBit := getBit(b[:], i) bBit := getBit(b[:], i)
oBit := getBit(other[:], i) oBit := getBit(other[:], i)
setBit(ret[:], i, bBit != oBit != carry) setBit(ret[:], i, bBit != oBit != carry)
@ -161,7 +161,7 @@ func (b Bitmap) add(other Bitmap) (Bitmap, bool) {
func (b Bitmap) Add(other Bitmap) Bitmap { func (b Bitmap) Add(other Bitmap) Bitmap {
ret, carry := b.add(other) ret, carry := b.add(other)
if carry { if carry {
panic("overflow in bitmap addition. limited to " + strconv.Itoa(nodeIDBits) + " bits.") panic("overflow in bitmap addition. limited to " + strconv.Itoa(NumBits) + " bits.")
} }
return ret return ret
} }
@ -173,7 +173,7 @@ func (b Bitmap) Sub(other Bitmap) Bitmap {
// ToDo: Why is this not supported? Should it say not implemented? BitMap might have a generic use case outside of dht. // ToDo: Why is this not supported? Should it say not implemented? BitMap might have a generic use case outside of dht.
panic("negative bitmaps not supported") panic("negative bitmaps not supported")
} }
complement, _ := other.Not().add(BitmapFromShortHexP("1")) complement, _ := other.Not().add(FromShortHexP("1"))
ret, _ := b.add(complement) ret, _ := b.add(complement)
return ret return ret
} }
@ -199,7 +199,7 @@ func (b Bitmap) PrefixLen() int {
} }
} }
} }
return nodeIDBits return NumBits
} }
// Prefix returns a copy of b with the first n bits set to 1 (if `one` is true) or 0 (if `one` is false) // Prefix returns a copy of b with the first n bits set to 1 (if `one` is true) or 0 (if `one` is false)
@ -233,7 +233,7 @@ func (b Bitmap) Suffix(n int, one bool) Bitmap {
Outer: Outer:
for i := len(ret) - 1; i >= 0; i-- { for i := len(ret) - 1; i >= 0; i-- {
for j := 7; j >= 0; j-- { for j := 7; j >= 0; j-- {
if i*8+j >= nodeIDBits-n { if i*8+j >= NumBits-n {
if one { if one {
ret[i] |= 1 << uint(7-j) ret[i] |= 1 << uint(7-j)
} else { } else {
@ -261,15 +261,15 @@ func (b *Bitmap) UnmarshalBencode(encoded []byte) error {
if err != nil { if err != nil {
return err return err
} }
if len(str) != nodeIDLength { if len(str) != NumBytes {
return errors.Err("invalid bitmap length") return errors.Err("invalid bitmap length")
} }
copy(b[:], str) copy(b[:], str)
return nil return nil
} }
// BitmapFromBytes returns a bitmap as long as the byte array is of a specific length specified in the parameters. // FromBytes returns a bitmap as long as the byte array is of a specific length specified in the parameters.
func BitmapFromBytes(data []byte) (Bitmap, error) { func FromBytes(data []byte) (Bitmap, error) {
var bmp Bitmap var bmp Bitmap
if len(data) != len(bmp) { if len(data) != len(bmp) {
@ -280,71 +280,71 @@ func BitmapFromBytes(data []byte) (Bitmap, error) {
return bmp, nil return bmp, nil
} }
// BitmapFromBytesP returns a bitmap as long as the byte array is of a specific length specified in the parameters // FromBytesP returns a bitmap as long as the byte array is of a specific length specified in the parameters
// otherwise it wil panic. // otherwise it wil panic.
func BitmapFromBytesP(data []byte) Bitmap { func FromBytesP(data []byte) Bitmap {
bmp, err := BitmapFromBytes(data) bmp, err := FromBytes(data)
if err != nil { if err != nil {
panic(err) panic(err)
} }
return bmp return bmp
} }
//BitmapFromString returns a bitmap by converting the string to bytes and creating from bytes as long as the byte array //FromString returns a bitmap by converting the string to bytes and creating from bytes as long as the byte array
// is of a specific length specified in the parameters // is of a specific length specified in the parameters
func BitmapFromString(data string) (Bitmap, error) { func FromString(data string) (Bitmap, error) {
return BitmapFromBytes([]byte(data)) return FromBytes([]byte(data))
} }
//BitmapFromStringP returns a bitmap by converting the string to bytes and creating from bytes as long as the byte array //FromStringP returns a bitmap by converting the string to bytes and creating from bytes as long as the byte array
// is of a specific length specified in the parameters otherwise it wil panic. // is of a specific length specified in the parameters otherwise it wil panic.
func BitmapFromStringP(data string) Bitmap { func FromStringP(data string) Bitmap {
bmp, err := BitmapFromString(data) bmp, err := FromString(data)
if err != nil { if err != nil {
panic(err) panic(err)
} }
return bmp return bmp
} }
//BitmapFromHex returns a bitmap by converting the hex string to bytes and creating from bytes as long as the byte array //FromHex returns a bitmap by converting the hex string to bytes and creating from bytes as long as the byte array
// is of a specific length specified in the parameters // is of a specific length specified in the parameters
func BitmapFromHex(hexStr string) (Bitmap, error) { func FromHex(hexStr string) (Bitmap, error) {
decoded, err := hex.DecodeString(hexStr) decoded, err := hex.DecodeString(hexStr)
if err != nil { if err != nil {
return Bitmap{}, errors.Err(err) return Bitmap{}, errors.Err(err)
} }
return BitmapFromBytes(decoded) return FromBytes(decoded)
} }
//BitmapFromHexP returns a bitmap by converting the hex string to bytes and creating from bytes as long as the byte array //FromHexP returns a bitmap by converting the hex string to bytes and creating from bytes as long as the byte array
// is of a specific length specified in the parameters otherwise it wil panic. // is of a specific length specified in the parameters otherwise it wil panic.
func BitmapFromHexP(hexStr string) Bitmap { func FromHexP(hexStr string) Bitmap {
bmp, err := BitmapFromHex(hexStr) bmp, err := FromHex(hexStr)
if err != nil { if err != nil {
panic(err) panic(err)
} }
return bmp return bmp
} }
//BitmapFromShortHex returns a bitmap by converting the hex string to bytes, adding the leading zeros prefix to the //FromShortHex returns a bitmap by converting the hex string to bytes, adding the leading zeros prefix to the
// hex string and creating from bytes as long as the byte array is of a specific length specified in the parameters // hex string and creating from bytes as long as the byte array is of a specific length specified in the parameters
func BitmapFromShortHex(hexStr string) (Bitmap, error) { func FromShortHex(hexStr string) (Bitmap, error) {
return BitmapFromHex(strings.Repeat("0", nodeIDLength*2-len(hexStr)) + hexStr) return FromHex(strings.Repeat("0", NumBytes*2-len(hexStr)) + hexStr)
} }
//BitmapFromShortHexP returns a bitmap by converting the hex string to bytes, adding the leading zeros prefix to the //FromShortHexP returns a bitmap by converting the hex string to bytes, adding the leading zeros prefix to the
// hex string and creating from bytes as long as the byte array is of a specific length specified in the parameters // hex string and creating from bytes as long as the byte array is of a specific length specified in the parameters
// otherwise it wil panic. // otherwise it wil panic.
func BitmapFromShortHexP(hexStr string) Bitmap { func FromShortHexP(hexStr string) Bitmap {
bmp, err := BitmapFromShortHex(hexStr) bmp, err := FromShortHex(hexStr)
if err != nil { if err != nil {
panic(err) panic(err)
} }
return bmp return bmp
} }
// RandomBitmapP generates a cryptographically random bitmap with the confines of the parameters specified. // Rand generates a cryptographically random bitmap with the confines of the parameters specified.
func RandomBitmapP() Bitmap { func Rand() Bitmap {
var id Bitmap var id Bitmap
_, err := rand.Read(id[:]) _, err := rand.Read(id[:])
if err != nil { if err != nil {
@ -353,11 +353,11 @@ func RandomBitmapP() Bitmap {
return id return id
} }
// RandomBitmapInRangeP generates a cryptographically random bitmap and while it is greater than the high threshold // RandInRangeP generates a cryptographically random bitmap and while it is greater than the high threshold
// bitmap will subtract the diff between high and low until it is no longer greater that the high. // bitmap will subtract the diff between high and low until it is no longer greater that the high.
func RandomBitmapInRangeP(low, high Bitmap) Bitmap { func RandInRangeP(low, high Bitmap) Bitmap {
diff := high.Sub(low) diff := high.Sub(low)
r := RandomBitmapP() r := Rand()
for r.Greater(diff) { for r.Greater(diff) {
r = r.Sub(diff) r = r.Sub(diff)
} }

View file

@ -1,4 +1,4 @@
package dht package bits
import ( import (
"fmt" "fmt"
@ -47,8 +47,8 @@ func TestBitmap(t *testing.T) {
} }
id := "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" id := "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
if BitmapFromHexP(id).Hex() != id { if FromHexP(id).Hex() != id {
t.Error(BitmapFromHexP(id).Hex()) t.Error(FromHexP(id).Hex())
} }
} }
@ -64,7 +64,7 @@ func TestBitmap_GetBit(t *testing.T) {
{bit: 380, expected: true, panic: false}, {bit: 380, expected: true, panic: false},
} }
b := BitmapFromShortHexP("a") b := FromShortHexP("a")
for _, test := range tt { for _, test := range tt {
actual := getBit(b[:], test.bit) actual := getBit(b[:], test.bit)
@ -90,8 +90,8 @@ func TestBitmap_SetBit(t *testing.T) {
} }
for _, test := range tt { for _, test := range tt {
expected := BitmapFromShortHexP(test.expected) expected := FromShortHexP(test.expected)
actual := BitmapFromShortHexP(test.hex) actual := FromShortHexP(test.hex)
if test.panic { if test.panic {
assertPanic(t, fmt.Sprintf("setting bit %d to %t", test.bit, test.one), func() { setBit(actual[:], test.bit, test.one) }) assertPanic(t, fmt.Sprintf("setting bit %d to %t", test.bit, test.one), func() { setBit(actual[:], test.bit, test.one) })
} else { } else {
@ -119,8 +119,8 @@ func TestBitmap_FromHexShort(t *testing.T) {
} }
for _, test := range tt { for _, test := range tt {
short := BitmapFromShortHexP(test.short) short := FromShortHexP(test.short)
long := BitmapFromHexP(test.long) long := FromHexP(test.long)
if !short.Equals(long) { if !short.Equals(long) {
t.Errorf("short hex %s: expected %s, got %s", test.short, long.Hex(), short.Hex()) t.Errorf("short hex %s: expected %s, got %s", test.short, long.Hex(), short.Hex())
} }
@ -128,7 +128,7 @@ func TestBitmap_FromHexShort(t *testing.T) {
} }
func TestBitmapMarshal(t *testing.T) { func TestBitmapMarshal(t *testing.T) {
b := BitmapFromStringP("123456789012345678901234567890123456789012345678") b := FromStringP("123456789012345678901234567890123456789012345678")
encoded, err := bencode.EncodeBytes(b) encoded, err := bencode.EncodeBytes(b)
if err != nil { if err != nil {
t.Error(err) t.Error(err)
@ -146,7 +146,7 @@ func TestBitmapMarshalEmbedded(t *testing.T) {
C int C int
}{ }{
A: "1", A: "1",
B: BitmapFromStringP("222222222222222222222222222222222222222222222222"), B: FromStringP("222222222222222222222222222222222222222222222222"),
C: 3, C: 3,
} }
@ -162,7 +162,7 @@ func TestBitmapMarshalEmbedded(t *testing.T) {
func TestBitmapMarshalEmbedded2(t *testing.T) { func TestBitmapMarshalEmbedded2(t *testing.T) {
encoded, err := bencode.EncodeBytes([]interface{}{ encoded, err := bencode.EncodeBytes([]interface{}{
BitmapFromStringP("333333333333333333333333333333333333333333333333"), FromStringP("333333333333333333333333333333333333333333333333"),
}) })
if err != nil { if err != nil {
t.Error(err) t.Error(err)
@ -189,7 +189,7 @@ func TestBitmap_PrefixLen(t *testing.T) {
} }
for _, test := range tt { for _, test := range tt {
len := BitmapFromHexP(test.hex).PrefixLen() len := FromHexP(test.hex).PrefixLen()
if len != test.len { if len != test.len {
t.Errorf("got prefix len %d; expected %d for %s", len, test.len, test.hex) t.Errorf("got prefix len %d; expected %d for %s", len, test.len, test.hex)
} }
@ -197,7 +197,7 @@ func TestBitmap_PrefixLen(t *testing.T) {
} }
func TestBitmap_Prefix(t *testing.T) { func TestBitmap_Prefix(t *testing.T) {
allOne := BitmapFromHexP("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff") allOne := FromHexP("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
zerosTT := []struct { zerosTT := []struct {
zeros int zeros int
@ -213,21 +213,21 @@ func TestBitmap_Prefix(t *testing.T) {
} }
for _, test := range zerosTT { for _, test := range zerosTT {
expected := BitmapFromHexP(test.expected) expected := FromHexP(test.expected)
actual := allOne.Prefix(test.zeros, false) actual := allOne.Prefix(test.zeros, false)
if !actual.Equals(expected) { if !actual.Equals(expected) {
t.Errorf("%d zeros: got %s; expected %s", test.zeros, actual.Hex(), expected.Hex()) t.Errorf("%d zeros: got %s; expected %s", test.zeros, actual.Hex(), expected.Hex())
} }
} }
for i := 0; i < nodeIDLength*8; i++ { for i := 0; i < NumBits; i++ {
b := allOne.Prefix(i, false) b := allOne.Prefix(i, false)
if b.PrefixLen() != i { if b.PrefixLen() != i {
t.Errorf("got prefix len %d; expected %d for %s", b.PrefixLen(), i, b.Hex()) t.Errorf("got prefix len %d; expected %d for %s", b.PrefixLen(), i, b.Hex())
} }
} }
allZero := BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000") allZero := FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
onesTT := []struct { onesTT := []struct {
ones int ones int
@ -243,7 +243,7 @@ func TestBitmap_Prefix(t *testing.T) {
} }
for _, test := range onesTT { for _, test := range onesTT {
expected := BitmapFromHexP(test.expected) expected := FromHexP(test.expected)
actual := allZero.Prefix(test.ones, true) actual := allZero.Prefix(test.ones, true)
if !actual.Equals(expected) { if !actual.Equals(expected) {
t.Errorf("%d ones: got %s; expected %s", test.ones, actual.Hex(), expected.Hex()) t.Errorf("%d ones: got %s; expected %s", test.ones, actual.Hex(), expected.Hex())
@ -252,7 +252,7 @@ func TestBitmap_Prefix(t *testing.T) {
} }
func TestBitmap_Suffix(t *testing.T) { func TestBitmap_Suffix(t *testing.T) {
allOne := BitmapFromHexP("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff") allOne := FromHexP("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
zerosTT := []struct { zerosTT := []struct {
zeros int zeros int
@ -268,21 +268,21 @@ func TestBitmap_Suffix(t *testing.T) {
} }
for _, test := range zerosTT { for _, test := range zerosTT {
expected := BitmapFromHexP(test.expected) expected := FromHexP(test.expected)
actual := allOne.Suffix(test.zeros, false) actual := allOne.Suffix(test.zeros, false)
if !actual.Equals(expected) { if !actual.Equals(expected) {
t.Errorf("%d zeros: got %s; expected %s", test.zeros, actual.Hex(), expected.Hex()) t.Errorf("%d zeros: got %s; expected %s", test.zeros, actual.Hex(), expected.Hex())
} }
} }
for i := 0; i < nodeIDLength*8; i++ { for i := 0; i < NumBits; i++ {
b := allOne.Prefix(i, false) b := allOne.Prefix(i, false)
if b.PrefixLen() != i { if b.PrefixLen() != i {
t.Errorf("got prefix len %d; expected %d for %s", b.PrefixLen(), i, b.Hex()) t.Errorf("got prefix len %d; expected %d for %s", b.PrefixLen(), i, b.Hex())
} }
} }
allZero := BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000") allZero := FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
onesTT := []struct { onesTT := []struct {
ones int ones int
@ -298,7 +298,7 @@ func TestBitmap_Suffix(t *testing.T) {
} }
for _, test := range onesTT { for _, test := range onesTT {
expected := BitmapFromHexP(test.expected) expected := FromHexP(test.expected)
actual := allZero.Suffix(test.ones, true) actual := allZero.Suffix(test.ones, true)
if !actual.Equals(expected) { if !actual.Equals(expected) {
t.Errorf("%d ones: got %s; expected %s", test.ones, actual.Hex(), expected.Hex()) t.Errorf("%d ones: got %s; expected %s", test.ones, actual.Hex(), expected.Hex())
@ -324,9 +324,9 @@ func TestBitmap_Add(t *testing.T) {
} }
for _, test := range tt { for _, test := range tt {
a := BitmapFromShortHexP(test.a) a := FromShortHexP(test.a)
b := BitmapFromShortHexP(test.b) b := FromShortHexP(test.b)
expected := BitmapFromShortHexP(test.sum) expected := FromShortHexP(test.sum)
if test.panic { if test.panic {
assertPanic(t, fmt.Sprintf("adding %s and %s", test.a, test.b), func() { a.Add(b) }) assertPanic(t, fmt.Sprintf("adding %s and %s", test.a, test.b), func() { a.Add(b) })
} else { } else {
@ -358,9 +358,9 @@ func TestBitmap_Sub(t *testing.T) {
} }
for _, test := range tt { for _, test := range tt {
a := BitmapFromShortHexP(test.a) a := FromShortHexP(test.a)
b := BitmapFromShortHexP(test.b) b := FromShortHexP(test.b)
expected := BitmapFromShortHexP(test.sum) expected := FromShortHexP(test.sum)
if test.panic { if test.panic {
assertPanic(t, fmt.Sprintf("subtracting %s - %s", test.a, test.b), func() { a.Sub(b) }) assertPanic(t, fmt.Sprintf("subtracting %s - %s", test.a, test.b), func() { a.Sub(b) })
} else { } else {
@ -371,3 +371,12 @@ func TestBitmap_Sub(t *testing.T) {
} }
} }
} }
func assertPanic(t *testing.T, text string, f func()) {
defer func() {
if r := recover(); r == nil {
t.Errorf("%s: did not panic as expected", text)
}
}()
f()
}

View file

@ -6,6 +6,7 @@ import (
"sync" "sync"
"time" "time"
"github.com/lbryio/reflector.go/dht/bits"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@ -22,11 +23,11 @@ type BootstrapNode struct {
nlock *sync.RWMutex nlock *sync.RWMutex
nodes []peer nodes []peer
nodeKeys map[Bitmap]int nodeKeys map[bits.Bitmap]int
} }
// NewBootstrapNode returns a BootstrapNode pointer. // NewBootstrapNode returns a BootstrapNode pointer.
func NewBootstrapNode(id Bitmap, initialPingInterval, rePingInterval time.Duration) *BootstrapNode { func NewBootstrapNode(id bits.Bitmap, initialPingInterval, rePingInterval time.Duration) *BootstrapNode {
b := &BootstrapNode{ b := &BootstrapNode{
Node: *NewNode(id), Node: *NewNode(id),
@ -35,7 +36,7 @@ func NewBootstrapNode(id Bitmap, initialPingInterval, rePingInterval time.Durati
nlock: &sync.RWMutex{}, nlock: &sync.RWMutex{},
nodes: make([]peer, 0), nodes: make([]peer, 0),
nodeKeys: make(map[Bitmap]int), nodeKeys: make(map[bits.Bitmap]int),
} }
b.requestHandler = b.handleRequest b.requestHandler = b.handleRequest

View file

@ -3,10 +3,12 @@ package dht
import ( import (
"net" "net"
"testing" "testing"
"github.com/lbryio/reflector.go/dht/bits"
) )
func TestBootstrapPing(t *testing.T) { func TestBootstrapPing(t *testing.T) {
b := NewBootstrapNode(RandomBitmapP(), 10, bootstrapDefaultRefreshDuration) b := NewBootstrapNode(bits.Rand(), 10, bootstrapDefaultRefreshDuration)
listener, err := net.ListenPacket(network, "127.0.0.1:54320") listener, err := net.ListenPacket(network, "127.0.0.1:54320")
if err != nil { if err != nil {

View file

@ -9,6 +9,7 @@ import (
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lbryio/lbry.go/stopOnce" "github.com/lbryio/lbry.go/stopOnce"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/spf13/cast" "github.com/spf13/cast"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
@ -26,8 +27,8 @@ const (
// TODO: alternatively, have a global Config for constants. at least that way tests can modify the values // TODO: alternatively, have a global Config for constants. at least that way tests can modify the values
alpha = 3 // this is the constant alpha in the spec alpha = 3 // this is the constant alpha in the spec
bucketSize = 8 // this is the constant k in the spec bucketSize = 8 // this is the constant k in the spec
nodeIDLength = 48 // bytes. this is the constant B in the spec nodeIDLength = bits.NumBytes // bytes. this is the constant B in the spec
nodeIDBits = nodeIDLength * 8 // number of bits in node ID nodeIDBits = bits.NumBits // number of bits in node ID
messageIDLength = 20 // bytes. messageIDLength = 20 // bytes.
udpRetry = 3 udpRetry = 3
@ -85,7 +86,7 @@ type DHT struct {
// lock for announced list // lock for announced list
lock *sync.RWMutex lock *sync.RWMutex
// list of bitmaps that need to be reannounced periodically // list of bitmaps that need to be reannounced periodically
announced map[Bitmap]bool announced map[bits.Bitmap]bool
} }
// New returns a DHT pointer. If config is nil, then config will be set to the default config. // New returns a DHT pointer. If config is nil, then config will be set to the default config.
@ -106,7 +107,7 @@ func New(config *Config) (*DHT, error) {
stop: stopOnce.New(), stop: stopOnce.New(),
joined: make(chan struct{}), joined: make(chan struct{}),
lock: &sync.RWMutex{}, lock: &sync.RWMutex{},
announced: make(map[Bitmap]bool), announced: make(map[bits.Bitmap]bool),
} }
return d, nil return d, nil
} }
@ -188,7 +189,7 @@ func (dht *DHT) Ping(addr string) error {
return err return err
} }
tmpNode := Contact{ID: RandomBitmapP(), IP: raddr.IP, Port: raddr.Port} tmpNode := Contact{ID: bits.Rand(), IP: raddr.IP, Port: raddr.Port}
res := dht.node.Send(tmpNode, Request{Method: pingMethod}) res := dht.node.Send(tmpNode, Request{Method: pingMethod})
if res == nil { if res == nil {
return errors.Err("no response from node %s", addr) return errors.Err("no response from node %s", addr)
@ -198,7 +199,7 @@ func (dht *DHT) Ping(addr string) error {
} }
// Get returns the list of nodes that have the blob for the given hash // Get returns the list of nodes that have the blob for the given hash
func (dht *DHT) Get(hash Bitmap) ([]Contact, error) { func (dht *DHT) Get(hash bits.Bitmap) ([]Contact, error) {
contacts, found, err := FindContacts(dht.node, hash, true, dht.stop.Ch()) contacts, found, err := FindContacts(dht.node, hash, true, dht.stop.Ch())
if err != nil { if err != nil {
return nil, err return nil, err
@ -211,7 +212,7 @@ func (dht *DHT) Get(hash Bitmap) ([]Contact, error) {
} }
// Announce announces to the DHT that this node has the blob for the given hash // Announce announces to the DHT that this node has the blob for the given hash
func (dht *DHT) Announce(hash Bitmap) error { func (dht *DHT) Announce(hash bits.Bitmap) error {
contacts, _, err := FindContacts(dht.node, hash, false, dht.stop.Ch()) contacts, _, err := FindContacts(dht.node, hash, false, dht.stop.Ch())
if err != nil { if err != nil {
return err return err
@ -251,7 +252,7 @@ func (dht *DHT) startReannouncer() {
dht.lock.RLock() dht.lock.RLock()
for h := range dht.announced { for h := range dht.announced {
dht.stop.Add(1) dht.stop.Add(1)
go func(bm Bitmap) { go func(bm bits.Bitmap) {
defer dht.stop.Done() defer dht.stop.Done()
err := dht.Announce(bm) err := dht.Announce(bm)
if err != nil { if err != nil {
@ -264,7 +265,7 @@ func (dht *DHT) startReannouncer() {
} }
} }
func (dht *DHT) storeOnNode(hash Bitmap, c Contact) { func (dht *DHT) storeOnNode(hash bits.Bitmap, c Contact) {
// self-store // self-store
if dht.contact.Equals(c) { if dht.contact.Equals(c) {
dht.node.Store(hash, c) dht.node.Store(hash, c)
@ -322,12 +323,16 @@ func (dht *DHT) PrintState() {
} }
} }
func (dht DHT) ID() bits.Bitmap {
return dht.contact.ID
}
func getContact(nodeID, addr string) (Contact, error) { func getContact(nodeID, addr string) (Contact, error) {
var c Contact var c Contact
if nodeID == "" { if nodeID == "" {
c.ID = RandomBitmapP() c.ID = bits.Rand()
} else { } else {
c.ID = BitmapFromHexP(nodeID) c.ID = bits.FromHexP(nodeID)
} }
ip, port, err := net.SplitHostPort(addr) ip, port, err := net.SplitHostPort(addr)

View file

@ -5,6 +5,8 @@ import (
"sync" "sync"
"testing" "testing"
"time" "time"
"github.com/lbryio/reflector.go/dht/bits"
) )
func TestNodeFinder_FindNodes(t *testing.T) { func TestNodeFinder_FindNodes(t *testing.T) {
@ -16,7 +18,7 @@ func TestNodeFinder_FindNodes(t *testing.T) {
bs.Shutdown() bs.Shutdown()
}() }()
contacts, found, err := FindContacts(dhts[2].node, RandomBitmapP(), false, nil) contacts, found, err := FindContacts(dhts[2].node, bits.Rand(), false, nil)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
@ -64,7 +66,7 @@ func TestNodeFinder_FindNodes_NoBootstrap(t *testing.T) {
} }
}() }()
_, _, err := FindContacts(dhts[2].node, RandomBitmapP(), false, nil) _, _, err := FindContacts(dhts[2].node, bits.Rand(), false, nil)
if err == nil { if err == nil {
t.Fatal("contact finder should have errored saying that there are no contacts in the routing table") t.Fatal("contact finder should have errored saying that there are no contacts in the routing table")
} }
@ -79,8 +81,8 @@ func TestNodeFinder_FindValue(t *testing.T) {
bs.Shutdown() bs.Shutdown()
}() }()
blobHashToFind := RandomBitmapP() blobHashToFind := bits.Rand()
nodeToFind := Contact{ID: RandomBitmapP(), IP: net.IPv4(1, 2, 3, 4), Port: 5678} nodeToFind := Contact{ID: bits.Rand(), IP: net.IPv4(1, 2, 3, 4), Port: 5678}
dhts[0].node.store.Upsert(blobHashToFind, nodeToFind) dhts[0].node.store.Upsert(blobHashToFind, nodeToFind)
contacts, found, err := FindContacts(dhts[2].node, blobHashToFind, true, nil) contacts, found, err := FindContacts(dhts[2].node, blobHashToFind, true, nil)
@ -113,9 +115,9 @@ func TestDHT_LargeDHT(t *testing.T) {
}() }()
wg := &sync.WaitGroup{} wg := &sync.WaitGroup{}
ids := make([]Bitmap, nodes) ids := make([]bits.Bitmap, nodes)
for i := range ids { for i := range ids {
ids[i] = RandomBitmapP() ids[i] = bits.Rand()
wg.Add(1) wg.Add(1)
go func(index int) { go func(index int) {
defer wg.Done() defer wg.Done()
@ -127,7 +129,7 @@ func TestDHT_LargeDHT(t *testing.T) {
wg.Wait() wg.Wait()
// check that each node is in at learst 1 other routing table // check that each node is in at learst 1 other routing table
rtCounts := make(map[Bitmap]int) rtCounts := make(map[bits.Bitmap]int)
for _, d := range dhts { for _, d := range dhts {
for _, d2 := range dhts { for _, d2 := range dhts {
if d.node.id.Equals(d2.node.id) { if d.node.id.Equals(d2.node.id) {
@ -149,7 +151,7 @@ func TestDHT_LargeDHT(t *testing.T) {
} }
// check that each ID is stored by at least 3 nodes // check that each ID is stored by at least 3 nodes
storeCounts := make(map[Bitmap]int) storeCounts := make(map[bits.Bitmap]int)
for _, d := range dhts { for _, d := range dhts {
for _, id := range ids { for _, id := range ids {
if len(d.node.store.Get(id)) > 0 { if len(d.node.store.Get(id)) > 0 {

View file

@ -8,6 +8,7 @@ import (
"strings" "strings"
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
"github.com/spf13/cast" "github.com/spf13/cast"
@ -83,9 +84,9 @@ func newMessageID() messageID {
// Request represents the structured request from one node to another. // Request represents the structured request from one node to another.
type Request struct { type Request struct {
ID messageID ID messageID
NodeID Bitmap NodeID bits.Bitmap
Method string Method string
Arg *Bitmap Arg *bits.Bitmap
StoreArgs *storeArgs StoreArgs *storeArgs
} }
@ -95,7 +96,7 @@ func (r Request) MarshalBencode() ([]byte, error) {
if r.StoreArgs != nil { if r.StoreArgs != nil {
args = r.StoreArgs args = r.StoreArgs
} else if r.Arg != nil { } else if r.Arg != nil {
args = []Bitmap{*r.Arg} args = []bits.Bitmap{*r.Arg}
} else { } else {
args = []string{} // request must always have keys 0-4, so we use an empty list for PING args = []string{} // request must always have keys 0-4, so we use an empty list for PING
} }
@ -112,7 +113,7 @@ func (r Request) MarshalBencode() ([]byte, error) {
func (r *Request) UnmarshalBencode(b []byte) error { func (r *Request) UnmarshalBencode(b []byte) error {
var raw struct { var raw struct {
ID messageID `bencode:"1"` ID messageID `bencode:"1"`
NodeID Bitmap `bencode:"2"` NodeID bits.Bitmap `bencode:"2"`
Method string `bencode:"3"` Method string `bencode:"3"`
Args bencode.RawMessage `bencode:"4"` Args bencode.RawMessage `bencode:"4"`
} }
@ -132,7 +133,7 @@ func (r *Request) UnmarshalBencode(b []byte) error {
return errors.Prefix("request unmarshal", err) return errors.Prefix("request unmarshal", err)
} }
} else if len(raw.Args) > 2 { // 2 because an empty list is `le` } else if len(raw.Args) > 2 { // 2 because an empty list is `le`
tmp := []Bitmap{} tmp := []bits.Bitmap{}
err = bencode.DecodeBytes(raw.Args, &tmp) err = bencode.DecodeBytes(raw.Args, &tmp)
if err != nil { if err != nil {
return errors.Prefix("request unmarshal", err) return errors.Prefix("request unmarshal", err)
@ -154,14 +155,14 @@ func (r Request) argsDebug() string {
type storeArgsValue struct { type storeArgsValue struct {
Token string `bencode:"token"` Token string `bencode:"token"`
LbryID Bitmap `bencode:"lbryid"` LbryID bits.Bitmap `bencode:"lbryid"`
Port int `bencode:"port"` Port int `bencode:"port"`
} }
type storeArgs struct { type storeArgs struct {
BlobHash Bitmap BlobHash bits.Bitmap
Value storeArgsValue Value storeArgsValue
NodeID Bitmap // original publisher id? I think this is getting fixed in the new dht stuff NodeID bits.Bitmap // original publisher id? I think this is getting fixed in the new dht stuff
SelfStore bool // this is an int on the wire SelfStore bool // this is an int on the wire
} }
@ -231,7 +232,7 @@ func (s *storeArgs) UnmarshalBencode(b []byte) error {
// Response represents the structured response one node returns to another. // Response represents the structured response one node returns to another.
type Response struct { type Response struct {
ID messageID ID messageID
NodeID Bitmap NodeID bits.Bitmap
Data string Data string
Contacts []Contact Contacts []Contact
FindValueKey string FindValueKey string
@ -308,7 +309,7 @@ func (r Response) MarshalBencode() ([]byte, error) {
func (r *Response) UnmarshalBencode(b []byte) error { func (r *Response) UnmarshalBencode(b []byte) error {
var raw struct { var raw struct {
ID messageID `bencode:"1"` ID messageID `bencode:"1"`
NodeID Bitmap `bencode:"2"` NodeID bits.Bitmap `bencode:"2"`
Data bencode.RawMessage `bencode:"3"` Data bencode.RawMessage `bencode:"3"`
} }
err := bencode.DecodeBytes(b, &raw) err := bencode.DecodeBytes(b, &raw)
@ -377,7 +378,7 @@ func (r *Response) UnmarshalBencode(b []byte) error {
// Error represents an error message that is returned from one node to another in communication. // Error represents an error message that is returned from one node to another in communication.
type Error struct { type Error struct {
ID messageID ID messageID
NodeID Bitmap NodeID bits.Bitmap
ExceptionType string ExceptionType string
Response []string Response []string
} }
@ -397,7 +398,7 @@ func (e Error) MarshalBencode() ([]byte, error) {
func (e *Error) UnmarshalBencode(b []byte) error { func (e *Error) UnmarshalBencode(b []byte) error {
var raw struct { var raw struct {
ID messageID `bencode:"1"` ID messageID `bencode:"1"`
NodeID Bitmap `bencode:"2"` NodeID bits.Bitmap `bencode:"2"`
ExceptionType string `bencode:"3"` ExceptionType string `bencode:"3"`
Args interface{} `bencode:"4"` Args interface{} `bencode:"4"`
} }

View file

@ -9,6 +9,7 @@ import (
"testing" "testing"
"github.com/davecgh/go-spew/spew" "github.com/davecgh/go-spew/spew"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
) )
@ -76,10 +77,10 @@ func TestBencodeDecodeStoreArgs(t *testing.T) {
func TestBencodeFindNodesResponse(t *testing.T) { func TestBencodeFindNodesResponse(t *testing.T) {
res := Response{ res := Response{
ID: newMessageID(), ID: newMessageID(),
NodeID: RandomBitmapP(), NodeID: bits.Rand(),
Contacts: []Contact{ Contacts: []Contact{
{ID: RandomBitmapP(), IP: net.IPv4(1, 2, 3, 4).To4(), Port: 5678}, {ID: bits.Rand(), IP: net.IPv4(1, 2, 3, 4).To4(), Port: 5678},
{ID: RandomBitmapP(), IP: net.IPv4(4, 3, 2, 1).To4(), Port: 8765}, {ID: bits.Rand(), IP: net.IPv4(4, 3, 2, 1).To4(), Port: 8765},
}, },
} }
@ -100,11 +101,11 @@ func TestBencodeFindNodesResponse(t *testing.T) {
func TestBencodeFindValueResponse(t *testing.T) { func TestBencodeFindValueResponse(t *testing.T) {
res := Response{ res := Response{
ID: newMessageID(), ID: newMessageID(),
NodeID: RandomBitmapP(), NodeID: bits.Rand(),
FindValueKey: RandomBitmapP().rawString(), FindValueKey: bits.Rand().String(),
Token: "arst", Token: "arst",
Contacts: []Contact{ Contacts: []Contact{
{ID: RandomBitmapP(), IP: net.IPv4(1, 2, 3, 4).To4(), Port: 5678}, {ID: bits.Rand(), IP: net.IPv4(1, 2, 3, 4).To4(), Port: 5678},
}, },
} }

View file

@ -11,6 +11,7 @@ import (
"github.com/lbryio/errors.go" "github.com/lbryio/errors.go"
"github.com/lbryio/lbry.go/stopOnce" "github.com/lbryio/lbry.go/stopOnce"
"github.com/lbryio/lbry.go/util" "github.com/lbryio/lbry.go/util"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/davecgh/go-spew/spew" "github.com/davecgh/go-spew/spew"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
@ -39,7 +40,7 @@ type RequestHandlerFunc func(addr *net.UDPAddr, request Request)
// Node is a type representation of a node on the network. // Node is a type representation of a node on the network.
type Node struct { type Node struct {
// the node's id // the node's id
id Bitmap id bits.Bitmap
// UDP connection for sending and receiving data // UDP connection for sending and receiving data
conn UDPConn conn UDPConn
// true if we've closed the connection on purpose // true if we've closed the connection on purpose
@ -64,7 +65,7 @@ type Node struct {
} }
// NewNode returns an initialized Node's pointer. // NewNode returns an initialized Node's pointer.
func NewNode(id Bitmap) *Node { func NewNode(id bits.Bitmap) *Node {
return &Node{ return &Node{
id: id, id: id,
rt: newRoutingTable(id), rt: newRoutingTable(id),
@ -270,7 +271,7 @@ func (n *Node) handleRequest(addr *net.UDPAddr, request Request) {
} }
if contacts := n.store.Get(*request.Arg); len(contacts) > 0 { if contacts := n.store.Get(*request.Arg); len(contacts) > 0 {
res.FindValueKey = request.Arg.rawString() res.FindValueKey = request.Arg.String()
res.Contacts = contacts res.Contacts = contacts
} else { } else {
res.Contacts = n.rt.GetClosest(*request.Arg, bucketSize) res.Contacts = n.rt.GetClosest(*request.Arg, bucketSize)
@ -446,6 +447,6 @@ func (n *Node) startRoutingTableGrooming() {
} }
// Store stores a node contact in the node's contact store. // Store stores a node contact in the node's contact store.
func (n *Node) Store(hash Bitmap, c Contact) { func (n *Node) Store(hash bits.Bitmap, c Contact) {
n.store.Upsert(hash, c) n.store.Upsert(hash, c)
} }

View file

@ -7,6 +7,7 @@ import (
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lbryio/lbry.go/stopOnce" "github.com/lbryio/lbry.go/stopOnce"
"github.com/lbryio/reflector.go/dht/bits"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@ -16,7 +17,7 @@ import (
type contactFinder struct { type contactFinder struct {
findValue bool // true if we're using findValue findValue bool // true if we're using findValue
target Bitmap target bits.Bitmap
node *Node node *Node
stop *stopOnce.Stopper stop *stopOnce.Stopper
@ -29,13 +30,13 @@ type contactFinder struct {
shortlistMutex *sync.Mutex shortlistMutex *sync.Mutex
shortlist []Contact shortlist []Contact
shortlistAdded map[Bitmap]bool shortlistAdded map[bits.Bitmap]bool
outstandingRequestsMutex *sync.RWMutex outstandingRequestsMutex *sync.RWMutex
outstandingRequests uint outstandingRequests uint
} }
func FindContacts(node *Node, target Bitmap, findValue bool, upstreamStop stopOnce.Chan) ([]Contact, bool, error) { func FindContacts(node *Node, target bits.Bitmap, findValue bool, upstreamStop stopOnce.Chan) ([]Contact, bool, error) {
cf := &contactFinder{ cf := &contactFinder{
node: node, node: node,
target: target, target: target,
@ -43,7 +44,7 @@ func FindContacts(node *Node, target Bitmap, findValue bool, upstreamStop stopOn
findValueMutex: &sync.Mutex{}, findValueMutex: &sync.Mutex{},
activeContactsMutex: &sync.Mutex{}, activeContactsMutex: &sync.Mutex{},
shortlistMutex: &sync.Mutex{}, shortlistMutex: &sync.Mutex{},
shortlistAdded: make(map[Bitmap]bool), shortlistAdded: make(map[bits.Bitmap]bool),
stop: stopOnce.New(), stop: stopOnce.New(),
outstandingRequestsMutex: &sync.RWMutex{}, outstandingRequestsMutex: &sync.RWMutex{},
} }
@ -259,7 +260,7 @@ func (cf *contactFinder) areRequestsOutstanding() bool {
return cf.outstandingRequests > 0 return cf.outstandingRequests > 0
} }
func sortInPlace(contacts []Contact, target Bitmap) { func sortInPlace(contacts []Contact, target bits.Bitmap) {
toSort := make([]sortedContact, len(contacts)) toSort := make([]sortedContact, len(contacts))
for i, n := range contacts { for i, n := range contacts {

View file

@ -5,12 +5,13 @@ import (
"testing" "testing"
"time" "time"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
) )
func TestPing(t *testing.T) { func TestPing(t *testing.T) {
dhtNodeID := RandomBitmapP() dhtNodeID := bits.Rand()
testNodeID := RandomBitmapP() testNodeID := bits.Rand()
conn := newTestUDPConn("127.0.0.1:21217") conn := newTestUDPConn("127.0.0.1:21217")
@ -30,7 +31,7 @@ func TestPing(t *testing.T) {
data, err := bencode.EncodeBytes(map[string]interface{}{ data, err := bencode.EncodeBytes(map[string]interface{}{
headerTypeField: requestType, headerTypeField: requestType,
headerMessageIDField: messageID, headerMessageIDField: messageID,
headerNodeIDField: testNodeID.rawString(), headerNodeIDField: testNodeID.String(),
headerPayloadField: "ping", headerPayloadField: "ping",
headerArgsField: []string{}, headerArgsField: []string{},
}) })
@ -86,7 +87,7 @@ func TestPing(t *testing.T) {
rNodeID, ok := response[headerNodeIDField].(string) rNodeID, ok := response[headerNodeIDField].(string)
if !ok { if !ok {
t.Error("node ID is not a string") t.Error("node ID is not a string")
} else if rNodeID != dhtNodeID.rawString() { } else if rNodeID != dhtNodeID.String() {
t.Error("unexpected node ID") t.Error("unexpected node ID")
} }
} }
@ -106,8 +107,8 @@ func TestPing(t *testing.T) {
} }
func TestStore(t *testing.T) { func TestStore(t *testing.T) {
dhtNodeID := RandomBitmapP() dhtNodeID := bits.Rand()
testNodeID := RandomBitmapP() testNodeID := bits.Rand()
conn := newTestUDPConn("127.0.0.1:21217") conn := newTestUDPConn("127.0.0.1:21217")
@ -123,7 +124,7 @@ func TestStore(t *testing.T) {
defer dht.Shutdown() defer dht.Shutdown()
messageID := newMessageID() messageID := newMessageID()
blobHashToStore := RandomBitmapP() blobHashToStore := bits.Rand()
storeRequest := Request{ storeRequest := Request{
ID: messageID, ID: messageID,
@ -176,7 +177,7 @@ func TestStore(t *testing.T) {
} }
} }
verifyResponse(t, response, messageID, dhtNodeID.rawString()) verifyResponse(t, response, messageID, dhtNodeID.String())
_, ok := response[headerPayloadField] _, ok := response[headerPayloadField]
if !ok { if !ok {
@ -204,8 +205,8 @@ func TestStore(t *testing.T) {
} }
func TestFindNode(t *testing.T) { func TestFindNode(t *testing.T) {
dhtNodeID := RandomBitmapP() dhtNodeID := bits.Rand()
testNodeID := RandomBitmapP() testNodeID := bits.Rand()
conn := newTestUDPConn("127.0.0.1:21217") conn := newTestUDPConn("127.0.0.1:21217")
@ -223,13 +224,13 @@ func TestFindNode(t *testing.T) {
nodesToInsert := 3 nodesToInsert := 3
var nodes []Contact var nodes []Contact
for i := 0; i < nodesToInsert; i++ { for i := 0; i < nodesToInsert; i++ {
n := Contact{ID: RandomBitmapP(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i} n := Contact{ID: bits.Rand(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i}
nodes = append(nodes, n) nodes = append(nodes, n)
dht.node.rt.Update(n) dht.node.rt.Update(n)
} }
messageID := newMessageID() messageID := newMessageID()
blobHashToFind := RandomBitmapP() blobHashToFind := bits.Rand()
request := Request{ request := Request{
ID: messageID, ID: messageID,
@ -257,7 +258,7 @@ func TestFindNode(t *testing.T) {
} }
} }
verifyResponse(t, response, messageID, dhtNodeID.rawString()) verifyResponse(t, response, messageID, dhtNodeID.String())
_, ok := response[headerPayloadField] _, ok := response[headerPayloadField]
if !ok { if !ok {
@ -273,8 +274,8 @@ func TestFindNode(t *testing.T) {
} }
func TestFindValueExisting(t *testing.T) { func TestFindValueExisting(t *testing.T) {
dhtNodeID := RandomBitmapP() dhtNodeID := bits.Rand()
testNodeID := RandomBitmapP() testNodeID := bits.Rand()
conn := newTestUDPConn("127.0.0.1:21217") conn := newTestUDPConn("127.0.0.1:21217")
@ -291,16 +292,16 @@ func TestFindValueExisting(t *testing.T) {
nodesToInsert := 3 nodesToInsert := 3
for i := 0; i < nodesToInsert; i++ { for i := 0; i < nodesToInsert; i++ {
n := Contact{ID: RandomBitmapP(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i} n := Contact{ID: bits.Rand(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i}
dht.node.rt.Update(n) dht.node.rt.Update(n)
} }
//data, _ := hex.DecodeString("64313a30693065313a3132303a7de8e57d34e316abbb5a8a8da50dcd1ad4c80e0f313a3234383a7ce1b831dec8689e44f80f547d2dea171f6a625e1a4ff6c6165e645f953103dabeb068a622203f859c6c64658fd3aa3b313a33393a66696e6456616c7565313a346c34383aa47624b8e7ee1e54df0c45e2eb858feb0b705bd2a78d8b739be31ba188f4bd6f56b371c51fecc5280d5fd26ba4168e966565") //data, _ := hex.DecodeString("64313a30693065313a3132303a7de8e57d34e316abbb5a8a8da50dcd1ad4c80e0f313a3234383a7ce1b831dec8689e44f80f547d2dea171f6a625e1a4ff6c6165e645f953103dabeb068a622203f859c6c64658fd3aa3b313a33393a66696e6456616c7565313a346c34383aa47624b8e7ee1e54df0c45e2eb858feb0b705bd2a78d8b739be31ba188f4bd6f56b371c51fecc5280d5fd26ba4168e966565")
messageID := newMessageID() messageID := newMessageID()
valueToFind := RandomBitmapP() valueToFind := bits.Rand()
nodeToFind := Contact{ID: RandomBitmapP(), IP: net.ParseIP("1.2.3.4"), Port: 1286} nodeToFind := Contact{ID: bits.Rand(), IP: net.ParseIP("1.2.3.4"), Port: 1286}
dht.node.store.Upsert(valueToFind, nodeToFind) dht.node.store.Upsert(valueToFind, nodeToFind)
dht.node.store.Upsert(valueToFind, nodeToFind) dht.node.store.Upsert(valueToFind, nodeToFind)
dht.node.store.Upsert(valueToFind, nodeToFind) dht.node.store.Upsert(valueToFind, nodeToFind)
@ -331,7 +332,7 @@ func TestFindValueExisting(t *testing.T) {
} }
} }
verifyResponse(t, response, messageID, dhtNodeID.rawString()) verifyResponse(t, response, messageID, dhtNodeID.String())
_, ok := response[headerPayloadField] _, ok := response[headerPayloadField]
if !ok { if !ok {
@ -343,7 +344,7 @@ func TestFindValueExisting(t *testing.T) {
t.Fatal("payload is not a dictionary") t.Fatal("payload is not a dictionary")
} }
compactContacts, ok := payload[valueToFind.rawString()] compactContacts, ok := payload[valueToFind.String()]
if !ok { if !ok {
t.Fatal("payload is missing key for search value") t.Fatal("payload is missing key for search value")
} }
@ -357,8 +358,8 @@ func TestFindValueExisting(t *testing.T) {
} }
func TestFindValueFallbackToFindNode(t *testing.T) { func TestFindValueFallbackToFindNode(t *testing.T) {
dhtNodeID := RandomBitmapP() dhtNodeID := bits.Rand()
testNodeID := RandomBitmapP() testNodeID := bits.Rand()
conn := newTestUDPConn("127.0.0.1:21217") conn := newTestUDPConn("127.0.0.1:21217")
@ -376,13 +377,13 @@ func TestFindValueFallbackToFindNode(t *testing.T) {
nodesToInsert := 3 nodesToInsert := 3
var nodes []Contact var nodes []Contact
for i := 0; i < nodesToInsert; i++ { for i := 0; i < nodesToInsert; i++ {
n := Contact{ID: RandomBitmapP(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i} n := Contact{ID: bits.Rand(), IP: net.ParseIP("127.0.0.1"), Port: 10000 + i}
nodes = append(nodes, n) nodes = append(nodes, n)
dht.node.rt.Update(n) dht.node.rt.Update(n)
} }
messageID := newMessageID() messageID := newMessageID()
valueToFind := RandomBitmapP() valueToFind := bits.Rand()
request := Request{ request := Request{
ID: messageID, ID: messageID,
@ -410,7 +411,7 @@ func TestFindValueFallbackToFindNode(t *testing.T) {
} }
} }
verifyResponse(t, response, messageID, dhtNodeID.rawString()) verifyResponse(t, response, messageID, dhtNodeID.String())
_, ok := response[headerPayloadField] _, ok := response[headerPayloadField]
if !ok { if !ok {

View file

@ -13,6 +13,7 @@ import (
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lbryio/lbry.go/stopOnce" "github.com/lbryio/lbry.go/stopOnce"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/lyoshenka/bencode" "github.com/lyoshenka/bencode"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
@ -25,7 +26,7 @@ import (
// Contact is a type representation of another node that a specific node is in communication with. // Contact is a type representation of another node that a specific node is in communication with.
type Contact struct { type Contact struct {
ID Bitmap ID bits.Bitmap
IP net.IP IP net.IP
Port int Port int
} }
@ -74,7 +75,7 @@ func (c *Contact) UnmarshalCompact(b []byte) error {
} }
c.IP = net.IPv4(b[0], b[1], b[2], b[3]).To4() c.IP = net.IPv4(b[0], b[1], b[2], b[3]).To4()
c.Port = int(uint16(b[5]) | uint16(b[4])<<8) c.Port = int(uint16(b[5]) | uint16(b[4])<<8)
c.ID = BitmapFromBytesP(b[6:]) c.ID = bits.FromBytesP(b[6:])
return nil return nil
} }
@ -120,7 +121,7 @@ func (c *Contact) UnmarshalBencode(b []byte) error {
type sortedContact struct { type sortedContact struct {
contact Contact contact Contact
xorDistanceToTarget Bitmap xorDistanceToTarget bits.Bitmap
} }
type byXorDistance []sortedContact type byXorDistance []sortedContact
@ -222,7 +223,7 @@ func (b *bucket) UpdateContact(c Contact, insertIfNew bool) {
} }
// FailContact marks a contact as having failed, and removes it if it failed too many times // FailContact marks a contact as having failed, and removes it if it failed too many times
func (b *bucket) FailContact(id Bitmap) { func (b *bucket) FailContact(id bits.Bitmap) {
b.lock.Lock() b.lock.Lock()
defer b.lock.Unlock() defer b.lock.Unlock()
i := find(id, b.peers) i := find(id, b.peers)
@ -233,7 +234,7 @@ func (b *bucket) FailContact(id Bitmap) {
} }
// find returns the contact in the bucket, or nil if the bucket does not contain the contact // find returns the contact in the bucket, or nil if the bucket does not contain the contact
func find(id Bitmap, peers []peer) int { func find(id bits.Bitmap, peers []peer) int {
for i := range peers { for i := range peers {
if peers[i].Contact.ID.Equals(id) { if peers[i].Contact.ID.Equals(id) {
return i return i
@ -250,11 +251,11 @@ func (b *bucket) NeedsRefresh(refreshInterval time.Duration) bool {
} }
type routingTable struct { type routingTable struct {
id Bitmap id bits.Bitmap
buckets [nodeIDBits]bucket buckets [nodeIDBits]bucket
} }
func newRoutingTable(id Bitmap) *routingTable { func newRoutingTable(id bits.Bitmap) *routingTable {
var rt routingTable var rt routingTable
rt.id = id rt.id = id
for i := range rt.buckets { for i := range rt.buckets {
@ -301,7 +302,7 @@ func (rt *routingTable) Fail(c Contact) {
// GetClosest returns the closest `limit` contacts from the routing table // GetClosest returns the closest `limit` contacts from the routing table
// It marks each bucket it accesses as having been accessed // It marks each bucket it accesses as having been accessed
func (rt *routingTable) GetClosest(target Bitmap, limit int) []Contact { func (rt *routingTable) GetClosest(target bits.Bitmap, limit int) []Contact {
var toSort []sortedContact var toSort []sortedContact
var bucketNum int var bucketNum int
@ -335,7 +336,7 @@ func (rt *routingTable) GetClosest(target Bitmap, limit int) []Contact {
return contacts return contacts
} }
func appendContacts(contacts []sortedContact, b bucket, target Bitmap) []sortedContact { func appendContacts(contacts []sortedContact, b bucket, target bits.Bitmap) []sortedContact {
for _, contact := range b.Contacts() { for _, contact := range b.Contacts() {
contacts = append(contacts, sortedContact{contact, contact.ID.Xor(target)}) contacts = append(contacts, sortedContact{contact, contact.ID.Xor(target)})
} }
@ -353,8 +354,8 @@ func (rt *routingTable) Count() int {
// Range is a structure that holds a min and max bitmaps. The range is used in bucket sizing. // Range is a structure that holds a min and max bitmaps. The range is used in bucket sizing.
type Range struct { type Range struct {
start Bitmap start bits.Bitmap
end Bitmap end bits.Bitmap
} }
// BucketRanges returns a slice of ranges, where the `start` of each range is the smallest id that can // BucketRanges returns a slice of ranges, where the `start` of each range is the smallest id that can
@ -370,22 +371,22 @@ func (rt *routingTable) BucketRanges() []Range {
return ranges return ranges
} }
func (rt *routingTable) bucketNumFor(target Bitmap) int { func (rt *routingTable) bucketNumFor(target bits.Bitmap) int {
if rt.id.Equals(target) { if rt.id.Equals(target) {
panic("routing table does not have a bucket for its own id") panic("routing table does not have a bucket for its own id")
} }
return nodeIDBits - 1 - target.Xor(rt.id).PrefixLen() return nodeIDBits - 1 - target.Xor(rt.id).PrefixLen()
} }
func (rt *routingTable) bucketFor(target Bitmap) *bucket { func (rt *routingTable) bucketFor(target bits.Bitmap) *bucket {
return &rt.buckets[rt.bucketNumFor(target)] return &rt.buckets[rt.bucketNumFor(target)]
} }
func (rt *routingTable) GetIDsForRefresh(refreshInterval time.Duration) []Bitmap { func (rt *routingTable) GetIDsForRefresh(refreshInterval time.Duration) []bits.Bitmap {
var bitmaps []Bitmap var bitmaps []bits.Bitmap
for i, bucket := range rt.buckets { for i, bucket := range rt.buckets {
if bucket.NeedsRefresh(refreshInterval) { if bucket.NeedsRefresh(refreshInterval) {
bitmaps = append(bitmaps, RandomBitmapP().Prefix(i, false)) bitmaps = append(bitmaps, bits.Rand().Prefix(i, false))
} }
} }
return bitmaps return bitmaps
@ -416,7 +417,7 @@ func (rt *routingTable) UnmarshalJSON(b []byte) error {
return err return err
} }
rt.id, err = BitmapFromHex(data.ID) rt.id, err = bits.FromHex(data.ID)
if err != nil { if err != nil {
return errors.Prefix("decoding ID", err) return errors.Prefix("decoding ID", err)
} }
@ -427,7 +428,7 @@ func (rt *routingTable) UnmarshalJSON(b []byte) error {
return errors.Err("decoding contact %s: wrong number of parts", s) return errors.Err("decoding contact %s: wrong number of parts", s)
} }
var c Contact var c Contact
c.ID, err = BitmapFromHex(parts[0]) c.ID, err = bits.FromHex(parts[0])
if err != nil { if err != nil {
return errors.Err("decoding contact %s: invalid ID: %s", s, err) return errors.Err("decoding contact %s: invalid ID: %s", s, err)
} }
@ -451,7 +452,7 @@ func RoutingTableRefresh(n *Node, refreshInterval time.Duration, upstreamStop st
for _, id := range n.rt.GetIDsForRefresh(refreshInterval) { for _, id := range n.rt.GetIDsForRefresh(refreshInterval) {
done.Add(1) done.Add(1)
go func(id Bitmap) { go func(id bits.Bitmap) {
defer done.Done() defer done.Done()
_, _, err := FindContacts(n, id, false, upstreamStop) _, _, err := FindContacts(n, id, false, upstreamStop)
if err != nil { if err != nil {

View file

@ -8,24 +8,25 @@ import (
"strings" "strings"
"testing" "testing"
"github.com/lbryio/reflector.go/dht/bits"
"github.com/sebdah/goldie" "github.com/sebdah/goldie"
) )
func TestRoutingTable_bucketFor(t *testing.T) { func TestRoutingTable_bucketFor(t *testing.T) {
rt := newRoutingTable(BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")) rt := newRoutingTable(bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"))
var tests = []struct { var tests = []struct {
id Bitmap id bits.Bitmap
expected int expected int
}{ }{
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"), 0}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001"), 0},
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000002"), 1}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000002"), 1},
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000003"), 1}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000003"), 1},
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004"), 2}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004"), 2},
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005"), 2}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005"), 2},
{BitmapFromHexP("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000f"), 3}, {bits.FromHexP("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000f"), 3},
{BitmapFromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000010"), 4}, {bits.FromHexP("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000010"), 4},
{BitmapFromHexP("F00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"), 383}, {bits.FromHexP("F00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"), 383},
{BitmapFromHexP("F0000000000000000000000000000000F0000000000000000000000000F0000000000000000000000000000000000000"), 383}, {bits.FromHexP("F0000000000000000000000000000000F0000000000000000000000000F0000000000000000000000000000000000000"), 383},
} }
for _, tt := range tests { for _, tt := range tests {
@ -37,14 +38,14 @@ func TestRoutingTable_bucketFor(t *testing.T) {
} }
func TestRoutingTable_GetClosest(t *testing.T) { func TestRoutingTable_GetClosest(t *testing.T) {
n1 := BitmapFromHexP("FFFFFFFF0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000") n1 := bits.FromHexP("FFFFFFFF0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
n2 := BitmapFromHexP("FFFFFFF00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000") n2 := bits.FromHexP("FFFFFFF00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
n3 := BitmapFromHexP("111111110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000") n3 := bits.FromHexP("111111110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
rt := newRoutingTable(n1) rt := newRoutingTable(n1)
rt.Update(Contact{n2, net.ParseIP("127.0.0.1"), 8001}) rt.Update(Contact{n2, net.ParseIP("127.0.0.1"), 8001})
rt.Update(Contact{n3, net.ParseIP("127.0.0.1"), 8002}) rt.Update(Contact{n3, net.ParseIP("127.0.0.1"), 8002})
contacts := rt.GetClosest(BitmapFromHexP("222222220000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"), 1) contacts := rt.GetClosest(bits.FromHexP("222222220000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"), 1)
if len(contacts) != 1 { if len(contacts) != 1 {
t.Fail() t.Fail()
return return
@ -68,7 +69,7 @@ func TestRoutingTable_GetClosest(t *testing.T) {
func TestCompactEncoding(t *testing.T) { func TestCompactEncoding(t *testing.T) {
c := Contact{ c := Contact{
ID: BitmapFromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41"), ID: bits.FromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41"),
IP: net.ParseIP("1.2.3.4"), IP: net.ParseIP("1.2.3.4"),
Port: int(55<<8 + 66), Port: int(55<<8 + 66),
} }
@ -144,13 +145,13 @@ func TestRoutingTable_MoveToBack(t *testing.T) {
} }
func TestRoutingTable_BucketRanges(t *testing.T) { func TestRoutingTable_BucketRanges(t *testing.T) {
id := BitmapFromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41") id := bits.FromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41")
ranges := newRoutingTable(id).BucketRanges() ranges := newRoutingTable(id).BucketRanges()
if !ranges[0].start.Equals(ranges[0].end) { if !ranges[0].start.Equals(ranges[0].end) {
t.Error("first bucket should only fit exactly one id") t.Error("first bucket should only fit exactly one id")
} }
for i := 0; i < 1000; i++ { for i := 0; i < 1000; i++ {
randID := RandomBitmapP() randID := bits.Rand()
found := -1 found := -1
for i, r := range ranges { for i, r := range ranges {
if r.start.LessOrEqual(randID) && r.end.GreaterOrEqual(randID) { if r.start.LessOrEqual(randID) && r.end.GreaterOrEqual(randID) {
@ -168,17 +169,17 @@ func TestRoutingTable_BucketRanges(t *testing.T) {
} }
func TestRoutingTable_Save(t *testing.T) { func TestRoutingTable_Save(t *testing.T) {
id := BitmapFromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41") id := bits.FromHexP("1c8aff71b99462464d9eeac639595ab99664be3482cb91a29d87467515c7d9158fe72aa1f1582dab07d8f8b5db277f41")
rt := newRoutingTable(id) rt := newRoutingTable(id)
ranges := rt.BucketRanges() ranges := rt.BucketRanges()
for i, r := range ranges { for i, r := range ranges {
for j := 0; j < bucketSize; j++ { for j := 0; j < bucketSize; j++ {
toAdd := r.start.Add(BitmapFromShortHexP(strconv.Itoa(j))) toAdd := r.start.Add(bits.FromShortHexP(strconv.Itoa(j)))
if toAdd.LessOrEqual(r.end) { if toAdd.LessOrEqual(r.end) {
rt.Update(Contact{ rt.Update(Contact{
ID: r.start.Add(BitmapFromShortHexP(strconv.Itoa(j))), ID: r.start.Add(bits.FromShortHexP(strconv.Itoa(j))),
IP: net.ParseIP("1.2.3." + strconv.Itoa(j)), IP: net.ParseIP("1.2.3." + strconv.Itoa(j)),
Port: 1 + i*bucketSize + j, Port: 1 + i*bucketSize + j,
}) })

View file

@ -1,36 +1,40 @@
package dht package dht
import "sync" import (
"sync"
"github.com/lbryio/reflector.go/dht/bits"
)
// TODO: expire stored data after tExpire time // TODO: expire stored data after tExpire time
type contactStore struct { type contactStore struct {
// map of blob hashes to (map of node IDs to bools) // map of blob hashes to (map of node IDs to bools)
hashes map[Bitmap]map[Bitmap]bool hashes map[bits.Bitmap]map[bits.Bitmap]bool
// stores the peers themselves, so they can be updated in one place // stores the peers themselves, so they can be updated in one place
contacts map[Bitmap]Contact contacts map[bits.Bitmap]Contact
lock sync.RWMutex lock sync.RWMutex
} }
func newStore() *contactStore { func newStore() *contactStore {
return &contactStore{ return &contactStore{
hashes: make(map[Bitmap]map[Bitmap]bool), hashes: make(map[bits.Bitmap]map[bits.Bitmap]bool),
contacts: make(map[Bitmap]Contact), contacts: make(map[bits.Bitmap]Contact),
} }
} }
func (s *contactStore) Upsert(blobHash Bitmap, contact Contact) { func (s *contactStore) Upsert(blobHash bits.Bitmap, contact Contact) {
s.lock.Lock() s.lock.Lock()
defer s.lock.Unlock() defer s.lock.Unlock()
if _, ok := s.hashes[blobHash]; !ok { if _, ok := s.hashes[blobHash]; !ok {
s.hashes[blobHash] = make(map[Bitmap]bool) s.hashes[blobHash] = make(map[bits.Bitmap]bool)
} }
s.hashes[blobHash][contact.ID] = true s.hashes[blobHash][contact.ID] = true
s.contacts[contact.ID] = contact s.contacts[contact.ID] = contact
} }
func (s *contactStore) Get(blobHash Bitmap) []Contact { func (s *contactStore) Get(blobHash bits.Bitmap) []Contact {
s.lock.RLock() s.lock.RLock()
defer s.lock.RUnlock() defer s.lock.RUnlock()

View file

@ -8,6 +8,7 @@ import (
"time" "time"
"github.com/lbryio/lbry.go/errors" "github.com/lbryio/lbry.go/errors"
"github.com/lbryio/reflector.go/dht/bits"
) )
var testingDHTIP = "127.0.0.1" var testingDHTIP = "127.0.0.1"
@ -21,7 +22,7 @@ func TestingCreateDHT(t *testing.T, numNodes int, bootstrap, concurrent bool) (*
if bootstrap { if bootstrap {
bootstrapAddress := testingDHTIP + ":" + strconv.Itoa(testingDHTFirstPort) bootstrapAddress := testingDHTIP + ":" + strconv.Itoa(testingDHTFirstPort)
seeds = []string{bootstrapAddress} seeds = []string{bootstrapAddress}
bootstrapNode = NewBootstrapNode(RandomBitmapP(), 0, bootstrapDefaultRefreshDuration) bootstrapNode = NewBootstrapNode(bits.Rand(), 0, bootstrapDefaultRefreshDuration)
listener, err := net.ListenPacket(network, bootstrapAddress) listener, err := net.ListenPacket(network, bootstrapAddress)
if err != nil { if err != nil {
panic(err) panic(err)
@ -39,7 +40,7 @@ func TestingCreateDHT(t *testing.T, numNodes int, bootstrap, concurrent bool) (*
dhts := make([]*DHT, numNodes) dhts := make([]*DHT, numNodes)
for i := 0; i < numNodes; i++ { for i := 0; i < numNodes; i++ {
dht, err := New(&Config{Address: testingDHTIP + ":" + strconv.Itoa(firstPort+i), NodeID: RandomBitmapP().Hex(), SeedNodes: seeds}) dht, err := New(&Config{Address: testingDHTIP + ":" + strconv.Itoa(firstPort+i), NodeID: bits.Rand().Hex(), SeedNodes: seeds})
if err != nil { if err != nil {
panic(err) panic(err)
} }
@ -225,7 +226,7 @@ func verifyContacts(t *testing.T, contacts []interface{}, nodes []Contact) {
continue continue
} }
for _, n := range nodes { for _, n := range nodes {
if n.ID.rawString() == id { if n.ID.String() == id {
currNode = n currNode = n
currNodeFound = true currNodeFound = true
foundNodes[id] = true foundNodes[id] = true
@ -305,12 +306,3 @@ func verifyCompactContacts(t *testing.T, contacts []interface{}, nodes []Contact
} }
} }
} }
func assertPanic(t *testing.T, text string, f func()) {
defer func() {
if r := recover(); r == nil {
t.Errorf("%s: did not panic as expected", text)
}
}()
f()
}

View file

@ -10,6 +10,7 @@ import (
"time" "time"
"github.com/lbryio/lbry.go/stopOnce" "github.com/lbryio/lbry.go/stopOnce"
"github.com/lbryio/reflector.go/dht/bits"
) )
type tokenManager struct { type tokenManager struct {
@ -46,15 +47,15 @@ func (tm *tokenManager) Stop() {
tm.stop.StopAndWait() tm.stop.StopAndWait()
} }
func (tm *tokenManager) Get(nodeID Bitmap, addr *net.UDPAddr) string { func (tm *tokenManager) Get(nodeID bits.Bitmap, addr *net.UDPAddr) string {
return genToken(tm.secret, nodeID, addr) return genToken(tm.secret, nodeID, addr)
} }
func (tm *tokenManager) Verify(token string, nodeID Bitmap, addr *net.UDPAddr) bool { func (tm *tokenManager) Verify(token string, nodeID bits.Bitmap, addr *net.UDPAddr) bool {
return token == genToken(tm.secret, nodeID, addr) || token == genToken(tm.prevSecret, nodeID, addr) return token == genToken(tm.secret, nodeID, addr) || token == genToken(tm.prevSecret, nodeID, addr)
} }
func genToken(secret []byte, nodeID Bitmap, addr *net.UDPAddr) string { func genToken(secret []byte, nodeID bits.Bitmap, addr *net.UDPAddr) string {
buf := bytes.Buffer{} buf := bytes.Buffer{}
buf.Write(nodeID[:]) buf.Write(nodeID[:])
buf.Write(addr.IP) buf.Write(addr.IP)