package dht import ( "math/rand" "net" "sync" "time" "github.com/lbryio/reflector.go/dht/bits" log "github.com/sirupsen/logrus" ) const ( bootstrapDefaultRefreshDuration = 15 * time.Minute ) // BootstrapNode is a configured node setup for testing. type BootstrapNode struct { Node initialPingInterval time.Duration checkInterval time.Duration nlock *sync.RWMutex nodes map[bits.Bitmap]*peer nodeIDs []bits.Bitmap // necessary for efficient random ID selection } // NewBootstrapNode returns a BootstrapNode pointer. func NewBootstrapNode(id bits.Bitmap, initialPingInterval, rePingInterval time.Duration) *BootstrapNode { b := &BootstrapNode{ Node: *NewNode(id), initialPingInterval: initialPingInterval, checkInterval: rePingInterval, nlock: &sync.RWMutex{}, nodes: make(map[bits.Bitmap]*peer), nodeIDs: make([]bits.Bitmap, 0), } b.requestHandler = b.handleRequest return b } // Add manually adds a contact func (b *BootstrapNode) Add(c Contact) { b.upsert(c) } // Connect connects to the given connection and starts any background threads necessary func (b *BootstrapNode) Connect(conn UDPConn) error { err := b.Node.Connect(conn) if err != nil { return err } log.Debugf("[%s] bootstrap: node connected", b.id.HexShort()) go func() { t := time.NewTicker(b.checkInterval / 5) for { select { case <-t.C: b.check() case <-b.stop.Ch(): return } } }() return nil } // upsert adds the contact to the list, or updates the lastPinged time func (b *BootstrapNode) upsert(c Contact) { b.nlock.Lock() defer b.nlock.Unlock() if node, exists := b.nodes[c.ID]; exists { log.Debugf("[%s] bootstrap: touching contact %s", b.id.HexShort(), node.Contact.ID.HexShort()) node.Touch() return } log.Debugf("[%s] bootstrap: adding new contact %s", b.id.HexShort(), c.ID.HexShort()) b.nodes[c.ID] = &peer{c, time.Now(), 0} b.nodeIDs = append(b.nodeIDs, c.ID) } // remove removes the contact from the list func (b *BootstrapNode) remove(c Contact) { b.nlock.Lock() defer b.nlock.Unlock() _, exists := b.nodes[c.ID] if !exists { return } log.Debugf("[%s] bootstrap: removing contact %s", b.id.HexShort(), c.ID.HexShort()) delete(b.nodes, c.ID) for i := range b.nodeIDs { if b.nodeIDs[i].Equals(c.ID) { b.nodeIDs = append(b.nodeIDs[:i], b.nodeIDs[i+1:]...) break } } } // get returns up to `limit` random contacts from the list func (b *BootstrapNode) get(limit int) []Contact { b.nlock.RLock() defer b.nlock.RUnlock() if len(b.nodes) < limit { limit = len(b.nodes) } ret := make([]Contact, limit) for i, k := range randKeys(len(b.nodeIDs))[:limit] { ret[i] = b.nodes[b.nodeIDs[k]].Contact } return ret } // ping pings a node. if the node responds, it is added to the list. otherwise, it is removed func (b *BootstrapNode) ping(c Contact) { log.Debugf("[%s] bootstrap: pinging %s", b.id.HexShort(), c.ID.HexShort()) b.stop.Add(1) defer b.stop.Done() resCh := b.SendAsync(c, Request{Method: pingMethod}) var res *Response select { case res = <-resCh: case <-b.stop.Ch(): return } if res != nil && res.Data == pingSuccessResponse { b.upsert(c) } else { b.remove(c) } } func (b *BootstrapNode) check() { b.nlock.RLock() defer b.nlock.RUnlock() for i := range b.nodes { if !b.nodes[i].ActiveInLast(b.checkInterval) { go b.ping(b.nodes[i].Contact) } } } // handleRequest handles the requests received from udp. func (b *BootstrapNode) handleRequest(addr *net.UDPAddr, request Request) { switch request.Method { case pingMethod: err := b.sendMessage(addr, Response{ID: request.ID, NodeID: b.id, Data: pingSuccessResponse}) if err != nil { log.Error("error sending response message - ", err) } case findNodeMethod: if request.Arg == nil { log.Errorln("request is missing arg") return } err := b.sendMessage(addr, Response{ ID: request.ID, NodeID: b.id, Contacts: b.get(bucketSize), }) if err != nil { log.Error("error sending 'findnodemethod' response message - ", err) } } go func() { b.nlock.RLock() _, exists := b.nodes[request.NodeID] b.nlock.RUnlock() if !exists { log.Debugf("[%s] bootstrap: queuing %s to ping", b.id.HexShort(), request.NodeID.HexShort()) <-time.After(b.initialPingInterval) b.nlock.RLock() _, exists = b.nodes[request.NodeID] b.nlock.RUnlock() if !exists { b.ping(Contact{ID: request.NodeID, IP: addr.IP, Port: addr.Port}) } } }() } func randKeys(max int) []int { keys := make([]int, max) for k := range keys { keys[k] = k } rand.Shuffle(max, func(i, j int) { keys[i], keys[j] = keys[j], keys[i] }) return keys }