store.go (2863B)
1 // Copyright 2011 Dmitry Chestnykh. All rights reserved. 2 // Use of this source code is governed by a MIT-style 3 // license that can be found in the LICENSE file. 4 5 package captcha 6 7 import ( 8 "sync" 9 "time" 10 ) 11 12 // Store an object implementing Store interface can be registered with SetCustomStore 13 // function to handle storage and retrieval of captcha ids and solutions for 14 // them, replacing the default memory store. 15 // 16 // It is the responsibility of an object to delete expired and used captchas 17 // when necessary (for example, the default memory store collects them in Set 18 // method after the certain amount of captchas has been stored.) 19 type Store interface { 20 // Set sets the digits for the captcha id. 21 Set(id string, answer string) 22 23 // Get returns stored digits for the captcha id. Clear indicates 24 // whether the captcha must be deleted from the store. 25 Get(id string, clear bool) (answer string, err error) 26 } 27 28 // expValue stores timestamp and id of captchas. It is used in the list inside 29 // memoryStore for indexing generated captchas by timestamp to enable garbage 30 // collection of expired captchas. 31 type idByTimeValue struct { 32 timestamp time.Time 33 id string 34 answer string 35 } 36 37 // memoryStore is an internal store for captcha ids and their values. 38 type memoryStore struct { 39 sync.RWMutex 40 digitsById map[string]idByTimeValue 41 // Number of items stored since last collection. 42 numStored int 43 // Number of saved items that triggers collection. 44 collectNum int 45 // Expiration time of captchas. 46 expiration time.Duration 47 } 48 49 // NewMemoryStore returns a new standard memory store for captchas with the 50 // given collection threshold and expiration time (duration). The returned 51 // store must be registered with SetCustomStore to replace the default one. 52 func NewMemoryStore(collectNum int, expiration time.Duration) Store { 53 s := new(memoryStore) 54 s.digitsById = make(map[string]idByTimeValue) 55 s.collectNum = collectNum 56 s.expiration = expiration 57 return s 58 } 59 60 func (s *memoryStore) Set(id string, answer string) { 61 s.Lock() 62 s.digitsById[id] = idByTimeValue{time.Now(), id, answer} 63 s.numStored++ 64 if s.numStored <= s.collectNum { 65 s.Unlock() 66 return 67 } 68 s.Unlock() 69 go s.collect() 70 } 71 72 func (s *memoryStore) Get(id string, clear bool) (string, error) { 73 if !clear { 74 // When we don't need to clear captcha, acquire read lock. 75 s.RLock() 76 defer s.RUnlock() 77 } else { 78 s.Lock() 79 defer s.Unlock() 80 } 81 el, ok := s.digitsById[id] 82 if el.timestamp.Add(s.expiration).Before(time.Now()) { 83 return "", ErrCaptchaExpired 84 } 85 if !ok { 86 return "", ErrNotFound 87 } 88 if clear { 89 delete(s.digitsById, id) 90 } 91 return el.answer, nil 92 } 93 94 func (s *memoryStore) collect() { 95 now := time.Now() 96 s.Lock() 97 defer s.Unlock() 98 s.numStored = 0 99 for _, ev := range s.digitsById { 100 if ev.timestamp.Add(s.expiration).Before(now) { 101 delete(s.digitsById, ev.id) 102 } 103 } 104 }