github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/p2p/security/blacklist.go (about)

     1  package security
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"sync"
     7  	"time"
     8  
     9  	cfg "github.com/bytom/bytom/config"
    10  	dbm "github.com/bytom/bytom/database/leveldb"
    11  )
    12  
    13  const (
    14  	defaultBanDuration = time.Hour * 1
    15  	blacklistKey       = "BlacklistPeers"
    16  )
    17  
    18  var (
    19  	ErrConnectBannedPeer = errors.New("connect banned peer")
    20  )
    21  
    22  type Blacklist struct {
    23  	peers map[string]time.Time
    24  	db    dbm.DB
    25  
    26  	mtx sync.Mutex
    27  }
    28  
    29  func NewBlacklist(config *cfg.Config) *Blacklist {
    30  	return &Blacklist{
    31  		peers: make(map[string]time.Time),
    32  		db:    dbm.NewDB("blacklist", config.DBBackend, config.DBDir()),
    33  	}
    34  }
    35  
    36  //AddPeer add peer to blacklist
    37  func (bl *Blacklist) AddPeer(ip string) error {
    38  	bl.mtx.Lock()
    39  	defer bl.mtx.Unlock()
    40  
    41  	// delete expired banned peers
    42  	for peer, banEnd := range bl.peers {
    43  		if time.Now().Before(banEnd) {
    44  			delete(bl.peers, peer)
    45  		}
    46  	}
    47  	// add banned peer
    48  	bl.peers[ip] = time.Now().Add(defaultBanDuration)
    49  	dataJSON, err := json.Marshal(bl.peers)
    50  	if err != nil {
    51  		return err
    52  	}
    53  
    54  	bl.db.Set([]byte(blacklistKey), dataJSON)
    55  	return nil
    56  }
    57  
    58  func (bl *Blacklist) delPeer(ip string) error {
    59  	delete(bl.peers, ip)
    60  	dataJson, err := json.Marshal(bl.peers)
    61  	if err != nil {
    62  		return err
    63  	}
    64  
    65  	bl.db.Set([]byte(blacklistKey), dataJson)
    66  	return nil
    67  }
    68  
    69  func (bl *Blacklist) DoFilter(ip string, pubKey string) error {
    70  	bl.mtx.Lock()
    71  	defer bl.mtx.Unlock()
    72  
    73  	if banEnd, ok := bl.peers[ip]; ok {
    74  		if time.Now().Before(banEnd) {
    75  			return ErrConnectBannedPeer
    76  		}
    77  
    78  		if err := bl.delPeer(ip); err != nil {
    79  			return err
    80  		}
    81  	}
    82  
    83  	return nil
    84  }
    85  
    86  // LoadPeers load banned peers from db
    87  func (bl *Blacklist) LoadPeers() error {
    88  	bl.mtx.Lock()
    89  	defer bl.mtx.Unlock()
    90  
    91  	if dataJSON := bl.db.Get([]byte(blacklistKey)); dataJSON != nil {
    92  		if err := json.Unmarshal(dataJSON, &bl.peers); err != nil {
    93  			return err
    94  		}
    95  	}
    96  
    97  	return nil
    98  }