mirror of
https://github.com/mpolden/echoip.git
synced 2025-06-29 14:17:50 +02:00
Implement response cache
This commit is contained in:
54
http/cache.go
Normal file
54
http/cache.go
Normal file
@ -0,0 +1,54 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"hash/fnv"
|
||||
"net"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Cache struct {
|
||||
capacity int
|
||||
mu sync.RWMutex
|
||||
entries map[uint64]*Response
|
||||
keys []uint64
|
||||
}
|
||||
|
||||
func NewCache(capacity int) *Cache {
|
||||
if capacity < 0 {
|
||||
capacity = 0
|
||||
}
|
||||
return &Cache{
|
||||
capacity: capacity,
|
||||
entries: make(map[uint64]*Response),
|
||||
keys: make([]uint64, 0, capacity),
|
||||
}
|
||||
}
|
||||
|
||||
func key(ip net.IP) uint64 {
|
||||
h := fnv.New64a()
|
||||
h.Write(ip)
|
||||
return h.Sum64()
|
||||
}
|
||||
|
||||
func (c *Cache) Set(ip net.IP, resp *Response) {
|
||||
if c.capacity == 0 {
|
||||
return
|
||||
}
|
||||
k := key(ip)
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if len(c.entries) == c.capacity && c.capacity > 0 {
|
||||
delete(c.entries, c.keys[0])
|
||||
c.keys = c.keys[1:]
|
||||
}
|
||||
c.entries[k] = resp
|
||||
c.keys = append(c.keys, k)
|
||||
}
|
||||
|
||||
func (c *Cache) Get(ip net.IP) (*Response, bool) {
|
||||
k := key(ip)
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
r, ok := c.entries[k]
|
||||
return r, ok
|
||||
}
|
41
http/cache_test.go
Normal file
41
http/cache_test.go
Normal file
@ -0,0 +1,41 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCacheCapacity(t *testing.T) {
|
||||
var tests = []struct {
|
||||
addCount, capacity, size int
|
||||
}{
|
||||
{1, 0, 0},
|
||||
{1, 2, 1},
|
||||
{2, 2, 2},
|
||||
{3, 2, 2},
|
||||
}
|
||||
for i, tt := range tests {
|
||||
c := NewCache(tt.capacity)
|
||||
var responses []*Response
|
||||
for i := 0; i < tt.addCount; i++ {
|
||||
ip := net.ParseIP(fmt.Sprintf("192.0.2.%d", i))
|
||||
r := &Response{IP: ip}
|
||||
responses = append(responses, r)
|
||||
c.Set(ip, r)
|
||||
}
|
||||
if got := len(c.entries); got != tt.size {
|
||||
t.Errorf("#%d: len(entries) = %d, want %d", i, got, tt.size)
|
||||
}
|
||||
if tt.capacity > 0 && tt.addCount > tt.capacity && tt.capacity == tt.size {
|
||||
lastAdded := responses[tt.addCount-1]
|
||||
if _, ok := c.Get(lastAdded.IP); !ok {
|
||||
t.Errorf("#%d: Get(%s) = (_, %t), want (_, %t)", i, lastAdded.IP.String(), ok, !ok)
|
||||
}
|
||||
firstAdded := responses[0]
|
||||
if _, ok := c.Get(firstAdded.IP); ok {
|
||||
t.Errorf("#%d: Get(%s) = (_, %t), want (_, %t)", i, firstAdded.IP.String(), ok, !ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
15
http/http.go
15
http/http.go
@ -27,6 +27,7 @@ type Server struct {
|
||||
IPHeaders []string
|
||||
LookupAddr func(net.IP) (string, error)
|
||||
LookupPort func(net.IP, uint64) error
|
||||
cache *Cache
|
||||
gr geo.Reader
|
||||
}
|
||||
|
||||
@ -51,8 +52,8 @@ type PortResponse struct {
|
||||
Reachable bool `json:"reachable"`
|
||||
}
|
||||
|
||||
func New(db geo.Reader) *Server {
|
||||
return &Server{gr: db}
|
||||
func New(db geo.Reader, cache *Cache) *Server {
|
||||
return &Server{cache: cache, gr: db}
|
||||
}
|
||||
|
||||
func ipFromForwardedForHeader(v string) string {
|
||||
@ -93,6 +94,10 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
response, ok := s.cache.Get(ip)
|
||||
if ok {
|
||||
return *response, nil
|
||||
}
|
||||
ipDecimal := iputil.ToDecimal(ip)
|
||||
country, _ := s.gr.Country(ip)
|
||||
city, _ := s.gr.City(ip)
|
||||
@ -111,7 +116,7 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
|
||||
parsed := useragent.Parse(userAgentRaw)
|
||||
userAgent = &parsed
|
||||
}
|
||||
return Response{
|
||||
response = &Response{
|
||||
IP: ip,
|
||||
IPDecimal: ipDecimal,
|
||||
Country: country.Name,
|
||||
@ -124,7 +129,9 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
|
||||
ASN: autonomousSystemNumber,
|
||||
ASNOrg: asn.AutonomousSystemOrganization,
|
||||
UserAgent: userAgent,
|
||||
}, nil
|
||||
}
|
||||
s.cache.Set(ip, response)
|
||||
return *response, nil
|
||||
}
|
||||
|
||||
func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
|
||||
|
@ -31,7 +31,7 @@ func (t *testDb) ASN(net.IP) (geo.ASN, error) {
|
||||
func (t *testDb) IsEmpty() bool { return false }
|
||||
|
||||
func testServer() *Server {
|
||||
return &Server{gr: &testDb{}, LookupAddr: lookupAddr, LookupPort: lookupPort}
|
||||
return &Server{cache: NewCache(100), gr: &testDb{}, LookupAddr: lookupAddr, LookupPort: lookupPort}
|
||||
}
|
||||
|
||||
func httpGet(url string, acceptMediaType string, userAgent string) (string, int, error) {
|
||||
|
Reference in New Issue
Block a user