a simple go tool vanity url server.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

db.go 5.2KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260
  1. package vain
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "net/http"
  6. "os"
  7. "strings"
  8. "sync"
  9. "time"
  10. verrors "mcquay.me/vain/errors"
  11. "mcquay.me/vain/metrics"
  12. )
  13. // NewMemDB returns a functional MemDB.
  14. func NewMemDB(p string) (*MemDB, error) {
  15. m := &MemDB{
  16. filename: p,
  17. Users: map[Email]User{},
  18. TokToEmail: map[Token]Email{},
  19. Packages: map[path]Package{},
  20. Namespaces: map[namespace]Email{},
  21. }
  22. f, err := os.Open(p)
  23. if err != nil {
  24. // file doesn't exist yet
  25. return m, nil
  26. }
  27. err = json.NewDecoder(f).Decode(m)
  28. return m, err
  29. }
  30. // MemDB implements an in-memory, and disk-backed database for a vain server.
  31. type MemDB struct {
  32. filename string
  33. l sync.RWMutex
  34. Users map[Email]User
  35. TokToEmail map[Token]Email
  36. Packages map[path]Package
  37. Namespaces map[namespace]Email
  38. }
  39. // NSForToken creates an entry namespaces with a relation to the token.
  40. func (m *MemDB) NSForToken(ns namespace, tok Token) error {
  41. m.l.Lock()
  42. defer m.l.Unlock()
  43. e, ok := m.TokToEmail[tok]
  44. if !ok {
  45. return verrors.HTTP{
  46. Message: fmt.Sprintf("User for token %q not found", tok),
  47. Code: http.StatusNotFound,
  48. }
  49. }
  50. if owner, ok := m.Namespaces[ns]; !ok {
  51. m.Namespaces[ns] = e
  52. } else {
  53. if m.Namespaces[ns] != owner {
  54. return verrors.HTTP{
  55. Message: fmt.Sprintf("not authorized against namespace %q", ns),
  56. Code: http.StatusUnauthorized,
  57. }
  58. }
  59. }
  60. return m.flush(m.filename)
  61. }
  62. // Package fetches the package associated with path.
  63. func (m *MemDB) Package(pth string) (Package, error) {
  64. m.l.RLock()
  65. defer m.l.RUnlock()
  66. pkg, ok := m.Packages[path(pth)]
  67. if ok {
  68. return pkg, nil
  69. }
  70. var longest Package
  71. for _, p := range m.Packages {
  72. if splitPathHasPrefix(strings.Split(pth, "/"), strings.Split(p.Path, "/")) {
  73. if len(p.Path) > len(longest.Path) {
  74. longest = p
  75. }
  76. }
  77. }
  78. var err error
  79. if longest.Path == "" {
  80. err = verrors.HTTP{
  81. Message: fmt.Sprintf("couldn't find package %q", pth),
  82. Code: http.StatusNotFound,
  83. }
  84. }
  85. return longest, err
  86. }
  87. // AddPackage adds p into packages table.
  88. func (m *MemDB) AddPackage(p Package) error {
  89. m.l.Lock()
  90. m.Packages[path(p.Path)] = p
  91. m.l.Unlock()
  92. return m.flush(m.filename)
  93. }
  94. // RemovePackage removes package with given path
  95. func (m *MemDB) RemovePackage(pth path) error {
  96. m.l.Lock()
  97. delete(m.Packages, pth)
  98. m.l.Unlock()
  99. return m.flush(m.filename)
  100. }
  101. // PackageExists tells if a package with path is in the database.
  102. func (m *MemDB) PackageExists(pth path) bool {
  103. m.l.RLock()
  104. _, ok := m.Packages[path(pth)]
  105. m.l.RUnlock()
  106. return ok
  107. }
  108. // Pkgs returns all packages from the database
  109. func (m *MemDB) Pkgs() []Package {
  110. ps := []Package{}
  111. m.l.RLock()
  112. for _, p := range m.Packages {
  113. ps = append(ps, p)
  114. }
  115. m.l.RUnlock()
  116. return ps
  117. }
  118. // Register adds email to the database, returning an error if there was one.
  119. func (m *MemDB) Register(e Email) (Token, error) {
  120. m.l.Lock()
  121. defer m.l.Unlock()
  122. if _, ok := m.Users[e]; ok {
  123. return "", verrors.HTTP{
  124. Message: fmt.Sprintf("duplicate email %q", e),
  125. Code: http.StatusConflict,
  126. }
  127. }
  128. tok := FreshToken()
  129. m.Users[e] = User{
  130. Email: e,
  131. token: tok,
  132. Requested: time.Now(),
  133. }
  134. m.TokToEmail[tok] = e
  135. return tok, m.flush(m.filename)
  136. }
  137. // Confirm modifies the user with the given token. Used on register confirmation.
  138. func (m *MemDB) Confirm(tok Token) (Token, error) {
  139. m.l.Lock()
  140. defer m.l.Unlock()
  141. e, ok := m.TokToEmail[tok]
  142. if !ok {
  143. return "", verrors.HTTP{
  144. Message: fmt.Sprintf("bad token: %s", tok),
  145. Code: http.StatusNotFound,
  146. }
  147. }
  148. delete(m.TokToEmail, tok)
  149. tok = FreshToken()
  150. u, ok := m.Users[e]
  151. if !ok {
  152. return "", verrors.HTTP{
  153. Message: fmt.Sprintf("inconsistent db; found email for token %q, but no user for email %q", tok, e),
  154. Code: http.StatusInternalServerError,
  155. }
  156. }
  157. u.token = tok
  158. m.Users[e] = u
  159. m.TokToEmail[tok] = e
  160. return tok, m.flush(m.filename)
  161. }
  162. // Forgot is used fetch a user's token. It implements rudimentary rate
  163. // limiting.
  164. func (m *MemDB) Forgot(e Email, window time.Duration) (Token, error) {
  165. m.l.Lock()
  166. defer m.l.Unlock()
  167. u, ok := m.Users[e]
  168. if !ok {
  169. return "", verrors.HTTP{
  170. Message: fmt.Sprintf("could not find email %q in db", e),
  171. Code: http.StatusNotFound,
  172. }
  173. }
  174. if u.Requested.After(time.Now()) {
  175. return "", verrors.HTTP{
  176. Message: fmt.Sprintf("rate limit hit for %q; try again in %0.2f mins", u.Email, u.Requested.Sub(time.Now()).Minutes()),
  177. Code: http.StatusTooManyRequests,
  178. }
  179. }
  180. return u.token, nil
  181. }
  182. // Sync takes a lock, and flushes the data to disk.
  183. func (m *MemDB) Sync() error {
  184. m.l.RLock()
  185. defer m.l.RUnlock()
  186. return m.flush(m.filename)
  187. }
  188. // flush writes to disk, but expects the user to have taken the lock.
  189. func (m *MemDB) flush(p string) error {
  190. defer metrics.DBTime("flush")()
  191. f, err := os.Create(p)
  192. if err != nil {
  193. return err
  194. }
  195. return json.NewEncoder(f).Encode(&m)
  196. }
  197. func (m *MemDB) addUser(e Email) (Token, error) {
  198. tok := FreshToken()
  199. m.l.Lock()
  200. m.Users[e] = User{
  201. Email: e,
  202. token: tok,
  203. Requested: time.Now(),
  204. }
  205. m.TokToEmail[tok] = e
  206. m.l.Unlock()
  207. return tok, m.flush(m.filename)
  208. }
  209. func (m *MemDB) user(e Email) (User, error) {
  210. m.l.Lock()
  211. u, ok := m.Users[e]
  212. m.l.Unlock()
  213. var err error
  214. if !ok {
  215. err = verrors.HTTP{
  216. Message: fmt.Sprintf("couldn't find user %q", e),
  217. Code: http.StatusNotFound,
  218. }
  219. }
  220. return u, err
  221. }