p2p/security/blacklist.go

   1 package security
   2
   3 import (
   4         "encoding/json"
   5         "errors"
   6         "sync"
   7         "time"
   8
   9         cfg "github.com/vapor/config"
  10         dbm "github.com/vapor/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         bl.peers[ip] = time.Now().Add(defaultBanDuration)
  42         dataJSON, err := json.Marshal(bl.peers)
  43         if err != nil {
  44                 return err
  45         }
  46
  47         bl.db.Set([]byte(blacklistKey), dataJSON)
  48         return nil
  49 }
  50
  51 func (bl *Blacklist) delPeer(ip string) error {
  52         delete(bl.peers, ip)
  53         dataJson, err := json.Marshal(bl.peers)
  54         if err != nil {
  55                 return err
  56         }
  57
  58         bl.db.Set([]byte(blacklistKey), dataJson)
  59         return nil
  60 }
  61
  62 func (bl *Blacklist) DoFilter(ip string, pubKey string) error {
  63         bl.mtx.Lock()
  64         defer bl.mtx.Unlock()
  65
  66         if banEnd, ok := bl.peers[ip]; ok {
  67                 if time.Now().Before(banEnd) {
  68                         return ErrConnectBannedPeer
  69                 }
  70
  71                 if err := bl.delPeer(ip); err != nil {
  72                         return err
  73                 }
  74         }
  75
  76         return nil
  77 }
  78
  79 // LoadPeers load banned peers from db
  80 func (bl *Blacklist) LoadPeers() error {
  81         bl.mtx.Lock()
  82         defer bl.mtx.Unlock()
  83
  84         if dataJSON := bl.db.Get([]byte(blacklistKey)); dataJSON != nil {
  85                 if err := json.Unmarshal(dataJSON, &bl.peers); err != nil {
  86                         return err
  87                 }
  88         }
  89
  90         return nil
  91 }