-
Notifications
You must be signed in to change notification settings - Fork 1
/
gooff.go
153 lines (129 loc) · 3.28 KB
/
gooff.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
package gooff
import (
"bufio"
"bytes"
"fmt"
"io/ioutil"
"log"
"net/http"
"os"
"github.com/dgraph-io/badger"
)
// Transport implements http.RoundTripper. When set as Transport of http.Client,
// it will store online requests for offline usage
type Transport struct {
Transport http.RoundTripper
preferDatabase bool
only200 bool
db *badger.DB
}
func init() {
GoOffline("", true, true)
}
// GoOffline returns the default transport, including whether to prefer db
// And sets up the database
func GoOffline(dbPath string, preferDatabase, only200 bool) {
db, err := setupDatabase(dbPath)
if err != nil {
log.Fatal(err)
}
http.DefaultTransport = &Transport{http.DefaultTransport, preferDatabase, only200, db}
}
func setupDatabase(dbPath string) (*badger.DB, error) {
if dbPath == "" {
dbPath = "/tmp"
}
os.MkdirAll(dbPath, os.ModePerm)
opts := badger.DefaultOptions
opts.Dir = dbPath + "/gooff"
opts.ValueDir = dbPath + "/gooff"
return badger.Open(opts)
}
// RoundTrip is the core part of this module and implements http.RoundTripper.
//
// If prefer database is true, always try and return value of request from
// database
//
// If there is data, return it
// If there is no data, then send the request anyway
//
// Store the result, linking it to the request
// If the request errors, then attempt to pull from database, otherwise return
// as is
func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
var resp *http.Response
var err error
if t.preferDatabase {
if resp, err = t.fetch(req); err == nil {
return resp, nil
}
}
resp, err = t.Transport.RoundTrip(req)
if err != nil {
if resp, err = t.fetch(req); err == nil {
return resp, nil
}
return nil, err
}
if !t.only200 || (t.only200 && resp.StatusCode == 200) {
if err = t.store(req, resp); err != nil {
return nil, err
}
}
return resp, err
}
func (t *Transport) fetch(req *http.Request) (*http.Response, error) {
var valCopy []byte
err := t.db.View(func(txn *badger.Txn) error {
bytesKey, err := key(req)
if err != nil {
return err
}
item, err := txn.Get(bytesKey)
if err != nil {
return err
}
err = item.Value(func(val []byte) error {
valCopy = append([]byte{}, val...)
return nil
})
return nil
})
if err != nil {
return nil, err
}
return http.ReadResponse(bufio.NewReader(bytes.NewReader(valCopy)), req)
}
func (t *Transport) store(req *http.Request, res *http.Response) error {
return t.db.Update(func(txn *badger.Txn) error {
// capture the bytes
bodyBytes, _ := ioutil.ReadAll(res.Body)
// must readd the bytes
res.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
var b bytes.Buffer
writer := bufio.NewWriter(&b)
if err := res.Write(writer); err != nil {
return err
}
writer.Flush()
// must readd again
res.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
bytesKey, err := key(req)
if err != nil {
return err
}
return txn.Set(bytesKey, b.Bytes())
})
}
func key(req *http.Request) ([]byte, error) {
key := fmt.Sprintf("%s:%s", req.Method, req.URL.String())
if req.Body != nil {
bodyBytes, err := ioutil.ReadAll(req.Body)
if err != nil {
return []byte{}, err
}
req.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))
key += ":" + string(bodyBytes)
}
return []byte(key), nil
}