increase rate limiter in test mode

- unit tests
- TODO: migrate to using portainer-cli repo

api/backup/import.go

@@ -0,0 +1,168 @@
+package backup
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"time"
+
+	"github.com/boltdb/bolt"
+	"github.com/sirupsen/logrus"
+)
+
+// TODO: use portainer-cli to import
+// source: https://github.com/portainer/portainer-cli/blob/master/util/database/import.go
+
+func ImportJsonToDatabase(jsonFilePath, portainerDbPath string) error {
+	if _, err := os.Stat(jsonFilePath); err != nil {
+		return fmt.Errorf("import file not found: %s: %s", jsonFilePath, err)
+	}
+	// if _, err := os.Stat(portainerDbPath); err == nil {
+	// 	return fmt.Errorf("ERROR: database file already exists: %s", portainerDbPath)
+	// }
+
+	return importJson(jsonFilePath, portainerDbPath)
+}
+
+func importJson(jsonFilePath, portainerDbPath string) error {
+	backup := make(map[string]interface{})
+
+	s, err := ioutil.ReadFile(jsonFilePath)
+	if err != nil {
+		return err
+	}
+	//err = json.Unmarshal([]byte(s), &backup)
+	d := json.NewDecoder(bytes.NewReader(s))
+	d.UseNumber()
+	if err = d.Decode(&backup); err != nil {
+		return err
+	}
+
+	connection, err := bolt.Open(portainerDbPath, 0600, &bolt.Options{Timeout: 2 * time.Second})
+	if err != nil {
+		return err
+	}
+	defer connection.Close()
+
+	return connection.Update(func(tx *bolt.Tx) error {
+		for bucketName, v := range backup {
+			logrus.WithField("bucketName", bucketName).Printf("CreateBucketIfNotExists")
+			_, err := tx.CreateBucketIfNotExists([]byte(bucketName))
+			if err != nil {
+				logrus.WithError(err).WithField("bucketName", bucketName).Printf("CreateBucketIfNotExists")
+				return err
+			}
+
+			bucket := tx.Bucket([]byte(bucketName))
+
+			switch bucketName {
+			case "version":
+				versions, ok := v.(map[string]interface{})
+				if !ok {
+					logrus.WithField("obj", v).Errorf("failed to cast %s map[string]interface{}", bucketName)
+				} else {
+					// TODO: test if those exist...
+					Put(bucketName, bucket, "DB_VERSION", versions["DB_VERSION"])
+					Put(bucketName, bucket, "INSTANCE_ID", versions["INSTANCE_ID"])
+				}
+			case "dockerhub":
+				Put(bucketName, bucket, "DOCKERHUB", v)
+
+				//obj, ok := v.([]map[string]string)
+				//if !ok {
+				//	logrus.WithField("obj", v).Errorf("failed to cast %s []map[string]interface{}", bucketName)
+				//} else {
+				//	Put(bucketName, bucket, "DOCKERHUB", obj)
+				//}
+			case "ssl":
+				obj, ok := v.(map[string]interface{})
+				if !ok {
+					logrus.WithField("obj", v).Errorf("failed to cast %s map[string]interface{}", bucketName)
+				} else {
+					Put(bucketName, bucket, "SSL", obj)
+				}
+			case "settings":
+				obj, ok := v.(map[string]interface{})
+				if !ok {
+					logrus.WithField("obj", v).Errorf("failed to cast %s map[string]interface{}", bucketName)
+				} else {
+					Put(bucketName, bucket, "SETTINGS", obj)
+				}
+			case "tunnel_server":
+				obj, ok := v.(map[string]interface{})
+				if !ok {
+					logrus.WithField("obj", v).Errorf("failed to cast %s map[string]interface{}", bucketName)
+				} else {
+					Put(bucketName, bucket, "INFO", obj)
+				}
+			default:
+				objlist, ok := v.([]interface{})
+				if !ok {
+					logrus.WithField("obj", v).Errorf("failed to cast %s []inferface{}", bucketName)
+				} else {
+					for _, ivalue := range objlist {
+						value, ok := ivalue.(map[string]interface{})
+						if !ok {
+							logrus.WithField("obj", value).Errorf("failed to cast %s map[string]interface{}", bucketName)
+						} else {
+							var ok bool
+							var id interface{}
+							switch bucketName {
+							case "endpoint_relations":
+								id, ok = value["EndpointID"] // TODO: need to make into an int, then do that weird stringification
+							default:
+								id, ok = value["Id"]
+							}
+							if !ok {
+								// endpoint_relations: EndpointID
+								logrus.WithField("obj", value).Errorf("No Id field:%s ", bucketName)
+								id = "error"
+							}
+							n, ok := id.(json.Number)
+							if !ok {
+								logrus.WithField("id", id).WithField("value", value).Errorf("failed to cast %s to int", bucketName)
+							} else {
+								key, err := n.Int64()
+								if err != nil {
+									logrus.WithError(err).WithField("id", id).WithField("key", key).WithField("value", value).Errorf("failed to cast %s to int", bucketName)
+								} else {
+									Put(bucketName, bucket, string(ConvertToKey(int(key))), value)
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+
+		return nil
+	})
+}
+
+// Honestly, I dunno why...
+func ConvertToKey(v int) []byte {
+	b := make([]byte, 8)
+	binary.BigEndian.PutUint64(b, uint64(v))
+	return b
+}
+
+func Put(bucketName string, bucket *bolt.Bucket, key string, object interface{}) error {
+	//logrus.WithField("bucketName", bucketName).WithField("key", key).WithField("object", object).Printf("Put")
+	data, err := json.Marshal(object)
+	if err != nil {
+		logrus.WithError(err).WithField("bucketName", bucketName).WithField("key", key).WithField("object", object).Errorf("failed marshal to json: (bucket: %s), (key: %s)", bucketName, key)
+
+		return err
+	}
+
+	err = bucket.Put([]byte(key), data)
+	if err != nil {
+		logrus.WithError(err).Errorf("failed Put into boltdb: (bucket: %s), (key: %s)", bucketName, key)
+
+		return err
+	}
+	return nil
+}

api/backup/seed.go

@@ -0,0 +1,63 @@
+package backup
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+
+	"github.com/pkg/errors"
+	"github.com/portainer/portainer/api/database/boltdb"
+	"github.com/portainer/portainer/api/dataservices"
+	"github.com/portainer/portainer/api/http/offlinegate"
+)
+
+// SeedStore seeds the store with provided JSON/map data, will trigger system shutdown, when finished.
+// NOTE: THIS WILL COMPLETELY OVERWRITE THE CURRENT STORE - only use this for testing.
+func SeedStore(storeData map[string]interface{}, filestorePath string, gate *offlinegate.OfflineGate, datastore dataservices.DataStore, shutdownTrigger context.CancelFunc) error {
+	// write storeData map to a temporary file
+	file, err := writeMapToFile(storeData)
+	if err != nil {
+		return err
+	}
+	defer os.Remove(file.Name())
+
+	unlock := gate.Lock()
+	defer unlock()
+
+	if err := datastore.Close(); err != nil {
+		return errors.Wrap(err, "Failed to stop db")
+	}
+
+	storePath := filepath.Join(filestorePath, boltdb.DatabaseFileName)
+	// TODO: use portainer-cli to import
+	if err := ImportJsonToDatabase(file.Name(), storePath); err != nil {
+		return errors.Wrap(err, "Unable to import JSON data to database")
+	}
+
+	shutdownTrigger()
+	return nil
+}
+
+// writeMapToFile writes a map to a temporary file and returns the file.
+func writeMapToFile(mapData map[string]interface{}) (*os.File, error) {
+	// map (json export) -> string
+	jsonData, err := json.Marshal(mapData)
+	if err != nil {
+		return nil, errors.Wrap(err, "Unable to marshal map to json")
+	}
+
+	// write string (json) to temporary file
+	file, err := os.CreateTemp("", "temp-db-export-json")
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	_, err = file.Write(jsonData)
+	if err != nil {
+		return nil, err
+	}
+
+	return file, nil
+}

api/backup/seed_test.go

@@ -0,0 +1,24 @@
+package backup
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_writeMapToFile(t *testing.T) {
+	is := assert.New(t)
+
+	data := map[string]interface{}{
+		"key1": "value1",
+		"key2": "value2",
+	}
+
+	f, err := writeMapToFile(data)
+	defer os.Remove(f.Name())
+
+	is.NoError(err)
+
+	is.NotNil(f)
+}

api/cmd/portainer/main_test.go

@@ -34,6 +34,8 @@ func Test_enableFeaturesFromFlags(t *testing.T) {
 		{"oPeN-amT", true},
 		{"fdo", true},
 		{"FDO", true},
+		{"Db-SeED", true},
+		{"DB-SEED", true},
 	}
 	for _, test := range tests {
 		t.Run(fmt.Sprintf("%s succeeds:%v", test.featureFlag, test.isSupported), func(t *testing.T) {

api/http/handler/backup/handler.go

@@ -39,6 +39,9 @@ func NewHandler(bouncer *security.RequestBouncer, dataStore dataservices.DataSto
 	h.Handle("/backup", bouncer.RestrictedAccess(adminAccess(httperror.LoggerHandler(h.backup)))).Methods(http.MethodPost)
 	h.Handle("/restore", bouncer.PublicAccess(httperror.LoggerHandler(h.restore))).Methods(http.MethodPost)
 
+	// db seed functionality is only available if specifically set as feature flag
+	h.Handle("/seed", bouncer.RestrictedAccess(adminAccess(httperror.LoggerHandler(h.seed)))).Methods(http.MethodPost)
+
 	return h
 }
 

api/http/handler/backup/seed.go

@@ -0,0 +1,59 @@
+package backup
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/pkg/errors"
+	httperror "github.com/portainer/libhttp/error"
+	"github.com/portainer/libhttp/request"
+	portainer "github.com/portainer/portainer/api"
+	operations "github.com/portainer/portainer/api/backup"
+)
+
+type seedPayload map[string]interface{}
+
+func (p *seedPayload) Validate(r *http.Request) error {
+	if p == nil {
+		return errors.New("Invalid seed data")
+	}
+	return nil
+}
+
+// @id Seed
+// @summary Overrwrites the current key-val database with provided payload - ONLY USE FOR TESTING.
+// @description  Overrwrites the current key-val database with provided payload - ONLY USE FOR TESTING; enable with `db-seed` feature flag.
+// @description **Access policy**: admin
+// @tags backup
+// @security ApiKeyAuth
+// @security jwt
+// @accept json
+// @param body body seedPayload true "The db seed payload - used to override current DB state"
+// @success 200 "Success"
+// @failure 400 "Invalid request"
+// @failure 404 "Not found"
+// @failure 500 "Server error"
+// @router /seed [post]
+func (h *Handler) seed(w http.ResponseWriter, r *http.Request) *httperror.HandlerError {
+	if !h.dataStore.Settings().IsFeatureFlagEnabled(portainer.FeatDBSeed) {
+		return &httperror.HandlerError{
+			StatusCode: http.StatusNotFound,
+			Message:    "Not found",
+			Err:        fmt.Errorf("feature flag '%s' not set", portainer.FeatDBSeed),
+		}
+	}
+
+	h.adminMonitor.Stop()
+	defer h.adminMonitor.Start()
+
+	var payload seedPayload
+	if err := request.DecodeAndValidateJSONPayload(r, &payload); err != nil {
+		return &httperror.HandlerError{StatusCode: http.StatusBadRequest, Message: "Invalid request payload", Err: err}
+	}
+
+	if err := operations.SeedStore(payload, h.filestorePath, h.gate, h.dataStore, h.shutdownTrigger); err != nil {
+		return &httperror.HandlerError{StatusCode: http.StatusInternalServerError, Message: "Failed to seed the store", Err: err}
+	}
+
+	return nil
+}

api/http/handler/handler.go

@@ -157,9 +157,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	switch {
 	case strings.HasPrefix(r.URL.Path, "/api/auth"):
 		http.StripPrefix("/api", h.AuthHandler).ServeHTTP(w, r)
-	case strings.HasPrefix(r.URL.Path, "/api/backup"):
-		http.StripPrefix("/api", h.BackupHandler).ServeHTTP(w, r)
-	case strings.HasPrefix(r.URL.Path, "/api/restore"):
+	case strings.HasPrefix(r.URL.Path, "/api/backup") || strings.HasPrefix(r.URL.Path, "/api/restore") || strings.HasPrefix(r.URL.Path, "/api/seed"):
 		http.StripPrefix("/api", h.BackupHandler).ServeHTTP(w, r)
 	case strings.HasPrefix(r.URL.Path, "/api/custom_templates"):
 		http.StripPrefix("/api", h.CustomTemplatesHandler).ServeHTTP(w, r)

api/http/server.go

@@ -107,6 +107,10 @@ func (server *Server) Start() error {
 	requestBouncer := security.NewRequestBouncer(server.DataStore, server.JWTService, server.APIKeyService)
 
 	rateLimiter := security.NewRateLimiter(10, 1*time.Second, 1*time.Hour)
+	if server.DataStore.Settings().IsFeatureFlagEnabled(portainer.FeatDBSeed) {
+		rateLimiter.Max = 1000
+		rateLimiter.BanDuration = 1 * time.Second
+	}
 	offlineGate := offlinegate.NewOfflineGate()
 
 	var authHandler = auth.NewHandler(requestBouncer, rateLimiter)

api/portainer.go

@@ -1352,12 +1352,14 @@ const (
 const (
 	FeatOpenAMT Feature = "open-amt"
 	FeatFDO     Feature = "fdo"
+	FeatDBSeed  Feature = "db-seed"
 )
 
 // List of supported features
 var SupportedFeatureFlags = []Feature{
 	FeatOpenAMT,
 	FeatFDO,
+	FeatDBSeed,
 }
 
 const (