package metadata import ( "bytes" "encoding/json" "errors" "github.com/dgraph-io/badger/v3" badgerOptions "github.com/dgraph-io/badger/v3/options" "github.com/minio/sha256-simd" "io" "io/ioutil" "net/http" "net/url" "runtime" "sort" "strings" "time" ) type BadgerCacheStore struct { handle *badger.DB gcTicker *time.Ticker closeChannel chan bool } func NewBadgerCacheStore(path string) (*BadgerCacheStore, error) { options := badger.DefaultOptions(path) options.SyncWrites = false options.NumVersionsToKeep = 1 options.Compression = badgerOptions.ZSTD options.ZSTDCompressionLevel = 1 db, err := badger.Open(options) if err != nil { return nil, err } store := &BadgerCacheStore{ handle: db, gcTicker: time.NewTicker(time.Minute * 5), closeChannel: make(chan bool), } go func() { defer store.gcTicker.Stop() for { select { case <-store.gcTicker.C: for store.handle.RunValueLogGC(0.6) == nil { } case <-store.closeChannel: return } } }() runtime.SetFinalizer(store, func(s *BadgerCacheStore) { s.closeChannel <- true s.handle.Close() }) return store, nil } func getRequestKey(r *http.Request) (out []byte) { key := r.Method + ":" + r.URL.String() + ":" var headers []string for k, v := range r.Header { hk := strings.ToLower(k) if hk == "user-agent" || hk == "x-fetched-at" { continue } headers = append(headers, hk+":"+strings.Join(v, ",")) } sort.SliceStable(headers, func(i, j int) bool { return strings.Compare(headers[i], headers[j]) < 0 }) key += strings.Join(headers, ";") hasher := sha256.New() hasher.Write([]byte(key)) out = hasher.Sum(out) return } type encodedRequest struct { Method string `json:"method"` URL string `json:"url"` Headers map[string][]string `json:"headers"` } func encodeRequest(r *http.Request) (out []byte) { value := encodedRequest{ Method: r.Method, URL: r.URL.String(), Headers: r.Header, } out, _ = json.Marshal(value) return } type encodedResponse struct { Request encodedRequest `json:"request"` Headers map[string][]string `json:"headers"` Status string `json:"status"` StatusCode int `json:"status_code"` Proto string `json:"proto"` Body []byte `json:"body"` } func encodeResponse(r *http.Response) (out []byte) { body, err := ioutil.ReadAll(r.Body) if err != nil { return nil } value := encodedResponse{ Request: encodedRequest{ Method: r.Request.Method, URL: r.Request.URL.String(), Headers: r.Request.Header, }, Headers: r.Header, Status: r.Status, StatusCode: r.StatusCode, Proto: r.Proto, Body: body, } out, _ = json.Marshal(value) return } func (s *BadgerCacheStore) Get(request *http.Request) (response *http.Response, err error) { key := getRequestKey(request) err = s.handle.View(func(txn *badger.Txn) error { item, err := txn.Get(key) if err != nil { return err } err = item.Value(func(val []byte) error { value := &encodedResponse{} err := json.Unmarshal(val, value) if err != nil { return err } requestUri, _ := url.Parse(value.Request.URL) response = &http.Response{ Request: &http.Request{ Method: value.Request.Method, URL: requestUri, Header: value.Request.Headers, }, Header: value.Headers, Status: value.Status, StatusCode: value.StatusCode, Proto: value.Proto, Body: io.NopCloser(bytes.NewReader(value.Body)), } return nil }) return err }) return } func (s *BadgerCacheStore) Set(request *http.Request, response *http.Response) (*http.Response, error) { defer response.Body.Close() key := getRequestKey(request) byteValue := encodeResponse(response) if len(byteValue) == 0 { return nil, errors.New("could not encode response") } err := s.handle.Update(func(txn *badger.Txn) error { return txn.Set(key, byteValue) }) if err != nil { return nil, err } return s.Get(request) }