+
+
package potency
+
+import (
+ "crypto/sha256"
+ "hash"
+ "io"
+)
+
+type bodyIntercept struct {
+ source io.ReadCloser
+ sha256 hash.Hash
+}
+
+func newBodyIntercept(source io.ReadCloser) *bodyIntercept {
+ return &bodyIntercept{
+ source: source,
+ sha256: sha256.New(),
+ }
+}
+
+func (bi *bodyIntercept) Read(p []byte) (int, error) {
+ numBytes, err := bi.source.Read(p)
+ bi.sha256.Write(p[:numBytes])
+
+ return numBytes, err
+}
+
+func (bi *bodyIntercept) Close() error {
+ return bi.source.Close()
+}
+
+
+
package potency
+
+import (
+ "bytes"
+ "crypto/sha256"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/gopatchy/jsrest"
+)
+
+type Potency struct {
+ handler http.Handler
+
+ lifetime time.Duration
+
+ cache map[string]*savedResult
+ cacheOldest *savedResult
+ cacheNewest *savedResult
+ cacheMu sync.RWMutex
+
+ inProgress map[string]bool
+ inProgressMu sync.Mutex
+}
+
+type savedResult struct {
+ key string
+
+ method string
+ url string
+ requestHeader http.Header
+ sha256 []byte
+
+ statusCode int
+ responseHeader http.Header
+ responseBody []byte
+
+ added time.Time
+ newer *savedResult
+}
+
+var (
+ ErrConflict = errors.New("conflict")
+ ErrMismatch = errors.New("idempotency mismatch")
+ ErrBodyMismatch = fmt.Errorf("request body mismatch: %w", ErrMismatch)
+ ErrMethodMismatch = fmt.Errorf("HTTP method mismatch: %w", ErrMismatch)
+ ErrURLMismatch = fmt.Errorf("URL mismatch: %w", ErrMismatch)
+ ErrHeaderMismatch = fmt.Errorf("Header mismatch: %w", ErrMismatch)
+ ErrInvalidKey = errors.New("invalid Idempotency-Key")
+
+ criticalHeaders = []string{
+ "Accept",
+ "Authorization",
+ }
+)
+
+func NewPotency(handler http.Handler) *Potency {
+ return &Potency{
+ handler: handler,
+ lifetime: 6 * time.Hour,
+ cache: map[string]*savedResult{},
+ inProgress: map[string]bool{},
+ }
+}
+
+func (p *Potency) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+ val := r.Header.Get("Idempotency-Key")
+ if val == "" {
+ p.handler.ServeHTTP(w, r)
+ return
+ }
+
+ err := p.serveHTTP(w, r, val)
+ if err != nil {
+ jsrest.WriteError(w, err)
+ }
+}
+
+func (p *Potency) SetLifetime(lifetime time.Duration) {
+ p.cacheMu.Lock()
+ defer p.cacheMu.Unlock()
+
+ p.lifetime = lifetime
+}
+
+func (p *Potency) NumCached() int {
+ p.cacheMu.RLock()
+ defer p.cacheMu.RUnlock()
+
+ return len(p.cache)
+}
+
+func (p *Potency) serveHTTP(w http.ResponseWriter, r *http.Request, val string) error {
+ if len(val) < 2 || !strings.HasPrefix(val, `"`) || !strings.HasSuffix(val, `"`) {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "%s (%w)", val, ErrInvalidKey)
+ }
+
+ key := val[1 : len(val)-1]
+
+ saved := p.read(key)
+
+ if saved != nil {
+ if r.Method != saved.method {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "%s (%w)", r.Method, ErrMethodMismatch)
+ }
+
+ if r.URL.String() != saved.url {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "%s (%w)", r.URL.String(), ErrURLMismatch)
+ }
+
+ for _, h := range criticalHeaders {
+ if saved.requestHeader.Get(h) != r.Header.Get(h) {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "%s: %s (%w)", h, r.Header.Get(h), ErrHeaderMismatch)
+ }
+ }
+
+ h := sha256.New()
+
+ _, err := io.Copy(h, r.Body)
+ if err != nil {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "hash request body failed (%w)", err)
+ }
+
+ sha256 := h.Sum(nil)
+ if !bytes.Equal(sha256, saved.sha256) {
+ return jsrest.Errorf(jsrest.ErrBadRequest, "%s vs %s (%w)", sha256, saved.sha256, ErrBodyMismatch)
+ }
+
+ for key, vals := range saved.responseHeader {
+ w.Header().Set(key, vals[0])
+ }
+
+ w.WriteHeader(saved.statusCode)
+ _, _ = w.Write(saved.responseBody)
+
+ return nil
+ }
+
+ // Store miss, proceed to normal execution with interception
+ err := p.lockKey(key)
+ if err != nil {
+ return jsrest.Errorf(jsrest.ErrConflict, "%s", key)
+ }
+
+ defer p.unlockKey(key)
+
+ requestHeader := http.Header{}
+ for _, h := range criticalHeaders {
+ requestHeader.Set(h, r.Header.Get(h))
+ }
+
+ bi := newBodyIntercept(r.Body)
+ r.Body = bi
+
+ rwi := newResponseWriterIntercept(w)
+ w = rwi
+
+ p.handler.ServeHTTP(w, r)
+
+ save := &savedResult{
+ key: key,
+
+ method: r.Method,
+ url: r.URL.String(),
+ requestHeader: requestHeader,
+ sha256: bi.sha256.Sum(nil),
+
+ statusCode: rwi.statusCode,
+ responseHeader: rwi.Header(),
+ responseBody: rwi.buf.Bytes(),
+ }
+
+ p.write(save)
+
+ return nil
+}
+
+func (p *Potency) lockKey(key string) error {
+ p.inProgressMu.Lock()
+ defer p.inProgressMu.Unlock()
+
+ if p.inProgress[key] {
+ return ErrConflict
+ }
+
+ p.inProgress[key] = true
+
+ return nil
+}
+
+func (p *Potency) unlockKey(key string) {
+ p.inProgressMu.Lock()
+ defer p.inProgressMu.Unlock()
+
+ delete(p.inProgress, key)
+}
+
+func (p *Potency) read(key string) *savedResult {
+ p.cacheMu.RLock()
+ defer p.cacheMu.RUnlock()
+
+ return p.cache[key]
+}
+
+func (p *Potency) write(sr *savedResult) {
+ p.cacheMu.Lock()
+ defer p.cacheMu.Unlock()
+
+ sr.added = time.Now()
+
+ p.cache[sr.key] = sr
+
+ if p.cacheNewest != nil {
+ p.cacheNewest.newer = sr
+ }
+
+ p.cacheNewest = sr
+
+ if p.cacheOldest == nil {
+ p.cacheOldest = sr
+ }
+
+ p.removeExpired()
+}
+
+func (p *Potency) removeExpired() {
+ cutoff := time.Now().Add(-1 * p.lifetime)
+
+ for iter := p.cacheOldest; iter != nil && iter.added.Before(cutoff); iter = iter.newer {
+ delete(p.cache, iter.key)
+ p.cacheOldest = iter
+ }
+}
+
+
+
package potency
+
+import (
+ "bytes"
+ "net/http"
+)
+
+type responseWriterIntercept struct {
+ dest http.ResponseWriter
+ buf bytes.Buffer
+ statusCode int
+}
+
+func newResponseWriterIntercept(dest http.ResponseWriter) *responseWriterIntercept {
+ return &responseWriterIntercept{
+ dest: dest,
+ buf: bytes.Buffer{},
+ statusCode: http.StatusOK,
+ }
+}
+
+func (rwi *responseWriterIntercept) Header() http.Header {
+ return rwi.dest.Header()
+}
+
+func (rwi *responseWriterIntercept) Write(data []byte) (int, error) {
+ rwi.buf.Write(data)
+ return rwi.dest.Write(data)
+}
+
+func (rwi *responseWriterIntercept) WriteHeader(statusCode int) {
+ rwi.statusCode = statusCode
+ rwi.dest.WriteHeader(statusCode)
+}
+
+
+