77 lines
1.3 KiB
Go
77 lines
1.3 KiB
Go
package reseed
|
|
|
|
import (
|
|
"io/ioutil"
|
|
"net"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
type Blacklist struct {
|
|
blacklist map[string]bool
|
|
m sync.RWMutex
|
|
}
|
|
|
|
func NewBlacklist() *Blacklist {
|
|
return &Blacklist{blacklist: make(map[string]bool), m: sync.RWMutex{}}
|
|
}
|
|
|
|
func (s *Blacklist) LoadFile(file string) error {
|
|
if file != "" {
|
|
if content, err := ioutil.ReadFile(file); err == nil {
|
|
for _, ip := range strings.Split(string(content), "\n") {
|
|
s.BlockIP(ip)
|
|
}
|
|
} else {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *Blacklist) BlockIP(ip string) {
|
|
s.m.Lock()
|
|
defer s.m.Unlock()
|
|
|
|
s.blacklist[ip] = true
|
|
}
|
|
|
|
func (s *Blacklist) isBlocked(ip string) bool {
|
|
s.m.RLock()
|
|
defer s.m.RUnlock()
|
|
|
|
blocked, found := s.blacklist[ip]
|
|
|
|
return found && blocked
|
|
}
|
|
|
|
type blacklistListener struct {
|
|
*net.TCPListener
|
|
blacklist *Blacklist
|
|
}
|
|
|
|
func (ln blacklistListener) Accept() (net.Conn, error) {
|
|
tc, err := ln.AcceptTCP()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
ip, _, err := net.SplitHostPort(tc.RemoteAddr().String())
|
|
if err != nil {
|
|
tc.Close()
|
|
return tc, err
|
|
}
|
|
|
|
if ln.blacklist.isBlocked(ip) {
|
|
tc.Close()
|
|
return tc, nil
|
|
}
|
|
|
|
return tc, err
|
|
}
|
|
|
|
func newBlacklistListener(ln net.Listener, bl *Blacklist) blacklistListener {
|
|
return blacklistListener{ln.(*net.TCPListener), bl}
|
|
}
|