add limiter

cmd/drone-autoscaler/main.go

@@ -17,6 +17,7 @@ import (
 	"github.com/drone/autoscaler/drivers/digitalocean"
 	"github.com/drone/autoscaler/drivers/hetznercloud"
 	"github.com/drone/autoscaler/engine"
+	"github.com/drone/autoscaler/limiter"
 	"github.com/drone/autoscaler/metrics"
 	"github.com/drone/autoscaler/server"
 	"github.com/drone/autoscaler/slack"
@@ -69,7 +70,7 @@ func main() {
 	if conf.Slack.Webhook != "" {
 		servers = slack.New(conf, servers)
 	}
-	// instruments the store with prometheus metrics.
+	servers = limiter.Limit(servers, conf.License)
 	servers = metrics.ServerCount(servers)
 	defer db.Close()
 

config/config.go

@@ -9,6 +9,7 @@ import "time"
 type (
 	// Config stores the configuration settings.
 	Config struct {
+		License  string
 		Interval time.Duration `default:"5m"`
 
 		Slack struct {

drivers/digitalocean/create_test.go

@@ -11,15 +11,11 @@ import (
 
 	"github.com/digitalocean/godo"
 	"github.com/drone/autoscaler"
-	"github.com/golang/mock/gomock"
 
 	"github.com/h2non/gock"
 )
 
 func TestCreate(t *testing.T) {
-	controller := gomock.NewController(t)
-	defer controller.Finish()
-
 	defer gock.Off()
 
 	gock.New("https://api.digitalocean.com").

drivers/digitalocean/destroy_test.go

@@ -17,9 +17,6 @@ import (
 )
 
 func TestDestroy(t *testing.T) {
-	controller := gomock.NewController(t)
-	defer controller.Finish()
-
 	defer gock.Off()
 
 	gock.New("https://api.digitalocean.com").

engine/planner.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/drone/autoscaler"
+	"github.com/drone/autoscaler/limiter"
 	"github.com/drone/drone-go/drone"
 
 	"github.com/dchest/uniuri"
@@ -111,6 +112,11 @@ func (p *planner) alloc(ctx context.Context, n int) error {
 		}
 
 		err := p.servers.Create(ctx, server)
+		if limiter.IsError(err) {
+			logger.Warn().Err(err).
+				Msg("cannot create server")
+			return err
+		}
 		if err != nil {
 			logger.Error().Err(err).
 				Msg("cannot create server")

limiter/error.go

@@ -0,0 +1,28 @@
+// Copyright 2018 Drone.IO Inc
+// Use of this software is governed by the Business Source License
+// that can be found in the LICENSE file.
+
+package limiter
+
+import "errors"
+
+var (
+	// Indicates the system has reached the limit on the
+	// number servers that it can provision under the
+	// current license.
+	errServerLimitExceeded = errors.New("Server limit exceeded")
+
+	// Indicates the license is expried. No new servers are
+	// provisioned until the license is renewed.
+	errLicenseExpired = errors.New("License expired")
+)
+
+// IsError returns true if the error is a Limit error.
+func IsError(err error) bool {
+	switch err {
+	case errServerLimitExceeded, errLicenseExpired:
+		return true
+	default:
+		return false
+	}
+}

limiter/error_test.go

@@ -0,0 +1,27 @@
+// Copyright 2018 Drone.IO Inc
+// Use of this software is governed by the Business Source License
+// that can be found in the LICENSE file.
+
+package limiter
+
+import (
+	"database/sql"
+	"testing"
+)
+
+func TestIsError(t *testing.T) {
+	var tests = []struct {
+		err error
+		res bool
+	}{
+		{nil, false},
+		{errLicenseExpired, true},
+		{errLicenseExpired, true},
+		{sql.ErrNoRows, false},
+	}
+	for _, test := range tests {
+		if got, want := IsError(test.err), test.res; got != want {
+			t.Errorf("Want IsError %v, got %v", want, got)
+		}
+	}
+}

limiter/license.go

@@ -0,0 +1,41 @@
+// Copyright 2018 Drone.IO Inc
+// Use of this software is governed by the Business Source License
+// that can be found in the LICENSE file.
+
+package limiter
+
+import (
+	"crypto"
+	"encoding/json"
+	"time"
+
+	"github.com/o1egl/paseto"
+)
+
+// License represents a software license key.
+type License struct {
+	Key string    `json:"key"`
+	Pro string    `json:"pro"`
+	Sub string    `json:"sub"`
+	Lim int       `json:"lim"`
+	Iss time.Time `json:"iat"`
+	Exp time.Time `json:"exp"`
+}
+
+// Expired returns true if the license is expired.
+func (l *License) Expired() bool {
+	return l.Exp.IsZero() == false && time.Now().After(l.Exp)
+}
+
+// ParseVerify parses and verifies the token, and returns
+// a License from the token payload.
+func ParseVerify(token string, publicKey crypto.PublicKey) (*License, error) {
+	var payload []byte
+	err := paseto.NewV2().Verify(token, publicKey, &payload, nil)
+	if err != nil {
+		return nil, err
+	}
+	out := new(License)
+	err = json.Unmarshal(payload, out)
+	return out, err
+}

limiter/license_test.go

@@ -0,0 +1,84 @@
+// Copyright 2018 Drone.IO Inc
+// Use of this software is governed by the Business Source License
+// that can be found in the LICENSE file.
+
+package limiter
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"crypto/rand"
+
+	"github.com/o1egl/paseto"
+	"golang.org/x/crypto/ed25519"
+)
+
+func TestParseVerify(t *testing.T) {
+	public, private, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Error(err)
+	}
+
+	a := &License{
+		Pro: "stripe",
+		Sub: "cus_CS7Nimer2KxGNp",
+		Lim: 25,
+		Exp: time.Now().UTC(),
+	}
+
+	data, _ := json.Marshal(a)
+	token, err := paseto.NewV2().Sign(private, data)
+	if err != nil {
+		t.Errorf(token)
+	}
+
+	b, err := ParseVerify(token, public)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if want, got := a.Sub, b.Sub; want != got {
+		t.Errorf("Want Sub %s, got %s", want, got)
+	}
+	if want, got := a.Pro, b.Pro; want != got {
+		t.Errorf("Want Pro %s, got %s", want, got)
+	}
+	if want, got := a.Lim, b.Lim; want != got {
+		t.Errorf("Want Lim %d, got %d", want, got)
+	}
+	if want, got := a.Exp, b.Exp; want != got {
+		t.Errorf("Want Exp %s, got %s", want, got)
+	}
+}
+
+func TestExpired(t *testing.T) {
+	tests := []struct {
+		Exp     time.Time
+		expired bool
+	}{
+		// zero value indicates no time limit
+		{
+			Exp:     time.Time{},
+			expired: false,
+		},
+		// one hour in the future
+		{
+			Exp:     time.Now().Add(time.Hour),
+			expired: false,
+		},
+		// one hour in the past
+		{
+			Exp:     time.Now().Add(-1 * time.Hour),
+			expired: true,
+		},
+	}
+
+	for _, test := range tests {
+		l := License{Exp: test.Exp}
+		if got, want := l.Expired(), test.expired; got != want {
+			t.Errorf("Want expired %v, got %v for %s", want, got, l.Exp)
+		}
+	}
+}

limiter/limit.go

@@ -0,0 +1,47 @@
+// Copyright 2018 Drone.IO Inc
+// Use of this software is governed by the Business Source License
+// that can be found in the LICENSE file.
+
+package limiter
+
+import (
+	"encoding/pem"
+
+	"github.com/drone/autoscaler"
+
+	"github.com/rs/zerolog/log"
+	"golang.org/x/crypto/ed25519"
+)
+
+var publicKey = []byte(`
+-----BEGIN PUBLIC KEY-----
+GB/hFnXEg63vDZ2W6mKFhLxZTuxMrlN/C/0iVZ2LfPQ=
+-----END PUBLIC KEY-----
+`)
+
+// Limit wraps the ServerStore to limit server creation
+// within the limitions of the license.
+func Limit(server autoscaler.ServerStore, token string) autoscaler.ServerStore {
+	if token == "" {
+		// if the token is empty the software is being
+		// used without a license. We assume this is for
+		// trial purposes and grant limited trial access.
+		return &limiter{server, &License{
+			Lim: 5,
+		}}
+	}
+	block, _ := pem.Decode(publicKey)
+	license, err := ParseVerify(token, ed25519.PublicKey(block.Bytes))
+	if err != nil {
+		panic(err)
+	}
+	log.Info().
+		Str("key", license.Key).
+		Str("pro", license.Pro).
+		Str("sub", license.Sub).
+		Int("lim", license.Lim).
+		Time("iat", license.Iss).
+		Time("exp", license.Exp).
+		Msg("license verified")
+	return &limiter{server, license}
+}