/
store.go
221 lines (185 loc) · 6.31 KB
/
store.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
package banman
import (
"bytes"
"encoding/binary"
"errors"
"fmt"
"net"
"time"
"github.com/Actinium-project/acmwallet/walletdb"
)
var (
// byteOrder is the preferred byte order in which we should write things
// to disk.
byteOrder = binary.BigEndian
// banStoreBucket is the top level bucket of the Store that will contain
// all relevant sub-buckets.
banStoreBucket = []byte("ban-store")
// banBucket is the main index in which we keep track of IP networks and
// their absolute expiration time.
//
// The key is the IP network host and the value is the absolute
// expiration time.
banBucket = []byte("ban-index")
// reasonBucket is an index in which we keep track of why an IP network
// was banned.
//
// The key is the IP network and the value is the Reason.
reasonBucket = []byte("reason-index")
// ErrCorruptedStore is an error returned when we attempt to locate any
// of the ban-related buckets in the database but are unable to.
ErrCorruptedStore = errors.New("corrupted ban store")
// ErrUnsupportedIP is an error returned when we attempt to parse an
// unsupported IP address type.
ErrUnsupportedIP = errors.New("unsupported IP type")
)
// Status gathers all of the details regarding an IP network's ban status.
type Status struct {
// Banned determines whether the IP network is currently banned.
Banned bool
// Reason is the reason for which the IP network was banned.
Reason Reason
// Expiration is the absolute time in which the ban will expire.
Expiration time.Time
}
// Store is the store responsible for maintaining records of banned IP networks.
// It uses IP networks, rather than single IP addresses, in order to coalesce
// multiple IP addresses that are likely to be correlated.
type Store interface {
// BanIPNet creates a ban record for the IP network within the store for
// the given duration. A reason can also be provided to note why the IP
// network is being banned. The record will exist until a call to Status
// is made after the ban expiration.
BanIPNet(*net.IPNet, Reason, time.Duration) error
// Status returns the ban status for a given IP network.
Status(*net.IPNet) (Status, error)
}
// NewStore returns a Store backed by a database.
func NewStore(db walletdb.DB) (Store, error) {
return newBanStore(db)
}
// banStore is a concrete implementation of the Store interface backed by a
// database.
type banStore struct {
db walletdb.DB
}
// A compile-time constraint to ensure banStore satisfies the Store interface.
var _ Store = (*banStore)(nil)
// newBanStore creates a concrete implementation of the Store interface backed
// by a database.
func newBanStore(db walletdb.DB) (*banStore, error) {
s := &banStore{db: db}
// We'll ensure the expected buckets are created upon initialization.
err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
banStore, err := tx.CreateTopLevelBucket(banStoreBucket)
if err != nil {
return err
}
_, err = banStore.CreateBucketIfNotExists(banBucket)
if err != nil {
return err
}
_, err = banStore.CreateBucketIfNotExists(reasonBucket)
return err
})
if err != nil && err != walletdb.ErrBucketExists {
return nil, err
}
return s, nil
}
// BanIPNet creates a ban record for the IP network within the store for the
// given duration. A reason can also be provided to note why the IP network is
// being banned. The record will exist until a call to Status is made after the
// ban expiration.
func (s *banStore) BanIPNet(ipNet *net.IPNet, reason Reason, duration time.Duration) error {
return walletdb.Update(s.db, func(tx walletdb.ReadWriteTx) error {
banStore := tx.ReadWriteBucket(banStoreBucket)
if banStore == nil {
return ErrCorruptedStore
}
banIndex := banStore.NestedReadWriteBucket(banBucket)
if banIndex == nil {
return ErrCorruptedStore
}
reasonIndex := banStore.NestedReadWriteBucket(reasonBucket)
if reasonIndex == nil {
return ErrCorruptedStore
}
var ipNetBuf bytes.Buffer
if err := encodeIPNet(&ipNetBuf, ipNet); err != nil {
return fmt.Errorf("unable to encode %v: %v", ipNet, err)
}
k := ipNetBuf.Bytes()
return addBannedIPNet(banIndex, reasonIndex, k, reason, duration)
})
}
// addBannedIPNet adds an entry to the ban store for the given IP network.
func addBannedIPNet(banIndex, reasonIndex walletdb.ReadWriteBucket,
ipNetKey []byte, reason Reason, duration time.Duration) error {
var v [8]byte
banExpiration := time.Now().Add(duration)
byteOrder.PutUint64(v[:], uint64(banExpiration.Unix()))
if err := banIndex.Put(ipNetKey, v[:]); err != nil {
return err
}
return reasonIndex.Put(ipNetKey, []byte{byte(reason)})
}
// Status returns the ban status for a given IP network.
func (s *banStore) Status(ipNet *net.IPNet) (Status, error) {
var banStatus Status
err := walletdb.Update(s.db, func(tx walletdb.ReadWriteTx) error {
banStore := tx.ReadWriteBucket(banStoreBucket)
if banStore == nil {
return ErrCorruptedStore
}
banIndex := banStore.NestedReadWriteBucket(banBucket)
if banIndex == nil {
return ErrCorruptedStore
}
reasonIndex := banStore.NestedReadWriteBucket(reasonBucket)
if reasonIndex == nil {
return ErrCorruptedStore
}
var ipNetBuf bytes.Buffer
if err := encodeIPNet(&ipNetBuf, ipNet); err != nil {
return fmt.Errorf("unable to encode %v: %v", ipNet, err)
}
k := ipNetBuf.Bytes()
status := fetchStatus(banIndex, reasonIndex, k)
// If the IP network's ban duration has expired, we can remove
// its entry from the store.
if !time.Now().Before(status.Expiration) {
return removeBannedIPNet(banIndex, reasonIndex, k)
}
banStatus = status
return nil
})
if err != nil {
return Status{}, err
}
return banStatus, nil
}
// fetchStatus retrieves the ban status of the given IP network.
func fetchStatus(banIndex, reasonIndex walletdb.ReadWriteBucket,
ipNetKey []byte) Status {
v := banIndex.Get(ipNetKey)
if v == nil {
return Status{}
}
reason := Reason(reasonIndex.Get(ipNetKey)[0])
banExpiration := time.Unix(int64(byteOrder.Uint64(v)), 0)
return Status{
Banned: true,
Reason: reason,
Expiration: banExpiration,
}
}
// removeBannedIPNet removes all references to a banned IP network within the
// ban store.
func removeBannedIPNet(banIndex, reasonIndex walletdb.ReadWriteBucket,
ipNetKey []byte) error {
if err := banIndex.Delete(ipNetKey); err != nil {
return err
}
return reasonIndex.Delete(ipNetKey)
}