Add support for limiting the number of syncs in the datastore

README.md

@@ -58,15 +58,12 @@ There are a few command line flags that can be used to change the behaviour of t
      `data`)
   - `-maxsize $size`: changes the maximum size of a sync (in bytes) that can be accepted by the API. The default is set
     to 512000.
+  - `-maxsyncs $num`: change the maximum number of syncs that can be stored. Default is 100000.
 
 ## Roadmap
 
 There are a few missing features that I would like to add.
 
-Hig priority:
- - **Enforce storage limits**: only the limit on the sync size is currently enforced: more checks must be implemented to
-   prevent abuse
-
 Medium priority:
  - **Rate limiting**: this can be partially enforced by a reverse proxy in front of the API server but would be nice to
    have it as part of the server itself

cmd/ubsserver/main.go

@@ -17,13 +17,21 @@ const (
 	apiVersion         = "1.1.13"
 	infoMessage        = "Powered by uBrowserSync"
 	defaultMaxSyncSize = 512000
+	defaultMaxSyncs    = 10000
+)
+
+const (
+	statusOnline   = 1
+	statusOffline  = 2
+	statusReadOnly = 3
 )
 
 var (
 	store         syncstore.Store
 	maxSyncSize   = defaultMaxSyncSize
+	maxSyncs      = defaultMaxSyncs
 	listen        string
-	serviceStatus = 1
+	serviceStatus = statusOnline
 	sidRe         = regexp.MustCompile("^[[:xdigit:]]{32}$")
 
 	invalidRequestError = syncstore.NewSyncError(
@@ -85,6 +93,10 @@ func ensureJSONRequest(w http.ResponseWriter, req *http.Request) bool {
 }
 
 func createSync(w http.ResponseWriter, req *http.Request) {
+	if serviceStatus != statusOnline {
+		sendJSONError(w, syncstore.NewSyncsForbiddenError)
+		return
+	}
 	if !ensureJSONRequest(w, req) {
 		return
 	}
@@ -104,6 +116,9 @@ func createSync(w http.ResponseWriter, req *http.Request) {
 		sendJSONError(w, err)
 		return
 	}
+	if maxSyncs > 0 && store.Count() >= maxSyncs {
+		serviceStatus = statusReadOnly
+	}
 	sendJSONOk(w, resp)
 }
 
@@ -139,7 +154,7 @@ func updateSync(syncId string, w http.ResponseWriter, req *http.Request) {
 		return
 	}
 	body := new(syncstore.UpdateReq)
-	req.Body = http.MaxBytesReader(w, req.Body, int64(10000 + maxSyncSize))
+	req.Body = http.MaxBytesReader(w, req.Body, int64(10000+maxSyncSize))
 	err := json.NewDecoder(req.Body).Decode(&body)
 	if err != nil {
 		sendJSONError(w, invalidRequestError)
@@ -227,6 +242,7 @@ func init() {
 	flag.StringVar(&listen, "listen", ":8090", "listen address and port")
 	flag.StringVar(&storeFlag, "store", "fs:data", "blob store driver")
 	flag.IntVar(&maxSyncSize, "maxsize", defaultMaxSyncSize, "maximum size of a sync in bytes")
+	flag.IntVar(&maxSyncs, "maxsyncs", defaultMaxSyncs, "maximum number of syncs")
 	flag.Parse()
 
 	storeFlagTokens := strings.Split(storeFlag, ":")
@@ -259,10 +275,10 @@ func main() {
 
 	log.Println("HTTP server listening on", listen)
 	server := &http.Server{
-		Addr: listen,
-		Handler: mux,
-		ReadTimeout: 15 * time.Second,
-		WriteTimeout: 15 * time.Second,
+		Addr:           listen,
+		Handler:        mux,
+		ReadTimeout:    15 * time.Second,
+		WriteTimeout:   15 * time.Second,
 		MaxHeaderBytes: 5000}
 	err := server.ListenAndServe()
 	log.Println("HTTP server terminated", err)

syncstore/driver_fs.go

@@ -3,13 +3,16 @@ package syncstore
 import (
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strings"
 )
 
 type FSStore struct {
 	path string
+	count int
 }
 
 func NewFSStore(path string) (*FSStore, error) {
@@ -23,9 +26,20 @@ func NewFSStore(path string) (*FSStore, error) {
 	}
 	var s = new(FSStore)
 	s.path = p
+	s.updateFileCount()
 	return s, nil
 }
 
+func (drv *FSStore) updateFileCount() {
+	drv.count = 0
+	filepath.Walk(drv.path, func(path string, info fs.FileInfo, err error) error {
+		if !info.IsDir() && strings.HasSuffix(info.Name(), ".json") {
+			drv.count++
+		}
+		return nil
+	})
+}
+
 func (drv *FSStore) storePath(id string) string {
 	return filepath.Join(drv.path, id[0:3], fmt.Sprintf("%s.json", id))
 }
@@ -37,9 +51,13 @@ func (drv *FSStore) RawSave(s *Blob) error {
 	}
 	fileName := drv.storePath(s.ID)
 	dirName := filepath.Dir(fileName)
-	err = os.Mkdir(dirName, 0700)
-	if err != nil && !os.IsExist(err) {
-		return err
+	_, err = os.Stat(fileName)
+	if err != nil {
+		err = os.Mkdir(dirName, 0700)
+		if err != nil && !os.IsExist(err) {
+			return err
+		}
+		drv.count++
 	}
 	f, err := os.CreateTemp(dirName, "tmp-"+s.ID+".*")
 	if err != nil {
@@ -78,3 +96,7 @@ func (drv *FSStore) Exists(id string) bool {
 	}
 	return (!st.IsDir()) && (st.Size() > 0)
 }
+
+func (drv *FSStore) Count() int {
+	return drv.count
+}

syncstore/driver_mem.go

@@ -24,3 +24,7 @@ func (drv *MemStore) Exists(id string) bool {
 	_, r := (*drv)[id]
 	return r
 }
+
+func (drv *MemStore) Count() int {
+	return len(*drv)
+}

syncstore/errors.go

@@ -57,4 +57,9 @@ var (
 		"RequiredDataNotFoundException",
 		"Unable to find required data",
 		http.StatusBadRequest)
+
+	NewSyncsForbiddenError = NewSyncError(
+		"NewSyncsForbiddenException",
+		"The service is not accepting new syncs",
+		http.StatusMethodNotAllowed)
 )

syncstore/store.go

@@ -12,6 +12,7 @@ type StoreDriver interface {
 	RawSave(s *Blob) error
 	RawLoad(id string) (*Blob, error)
 	Exists(id string) bool
+	Count() int
 }
 
 type Store struct {
@@ -106,3 +107,7 @@ func (store *Store) GetVersion(id string) (string, error) {
 	}
 	return r.Version, nil
 }
+
+func (store *Store) Count() int {
+	return store.drv.Count()
+}