Variable length encoding

cmd/quota-control/main.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/0xsequence/quotacontrol/proto"
+	"github.com/spf13/cobra"
+)
+
+func main() {
+	var rootCmd = &cobra.Command{Use: "quota-control"}
+	var accessKeyCmd = &cobra.Command{
+		Use:   "accesskey",
+		Short: "Manage access keys",
+	}
+
+	var verifyCmd = &cobra.Command{
+		Use:   "verify",
+		Short: "Verify an access key",
+		Run: func(cmd *cobra.Command, args []string) {
+			if len(args) != 1 {
+				fmt.Println("Usage: verify <access_key>")
+				os.Exit(1)
+			}
+			fmt.Println("Verifying access key...")
+			var errs []error
+			for i := len(proto.SupportedEncodings) - 1; i >= 0; i-- {
+				encoding := proto.SupportedEncodings[i]
+				v := encoding.Version()
+
+				projectID, ecosystemID, err := encoding.Decode(args[0])
+				if err != nil {
+					errs = append(errs, fmt.Errorf("v%d: %v", v, err))
+					continue
+				}
+				fmt.Printf("Access key decoded => version:%d, projectID:%d, ecosystemID:%d\n", v, projectID, ecosystemID)
+				return
+			}
+
+			fmt.Println("Access key is invalid:")
+			for _, err := range errs {
+				fmt.Println("-", err)
+			}
+		},
+	}
+
+	accessKeyCmd.AddCommand(verifyCmd)
+	rootCmd.AddCommand(accessKeyCmd)
+
+	if err := rootCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+}

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/jxskiss/base62 v1.1.0
 	github.com/redis/go-redis/v9 v9.7.0
+	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
 )
 
@@ -33,6 +34,7 @@ require (
 	github.com/goccy/go-json v0.10.3 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/holiman/uint256 v1.3.2 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/lestrrat-go/blackmagic v1.0.2 // indirect
 	github.com/lestrrat-go/httpcc v1.0.1 // indirect
@@ -42,6 +44,7 @@ require (
 	github.com/lestrrat-go/option v1.0.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/segmentio/asm v1.2.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/yuin/gopher-lua v1.1.1 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect

go.sum

@@ -14,6 +14,7 @@ github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
 github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
@@ -48,6 +49,8 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs
 github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/holiman/uint256 v1.3.2 h1:a9EgMPSC1AAaj1SZL5zIQD3WbwTuHrMGOerLjGmM/TA=
 github.com/holiman/uint256 v1.3.2/go.mod h1:EOMSn4q6Nyt9P6efbI3bueV4e1b3dGlUCXeiRV4ng7E=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jxskiss/base62 v1.1.0 h1:A5zbF8v8WXx2xixnAKD2w+abC+sIzYJX+nxmhA6HWFw=
 github.com/jxskiss/base62 v1.1.0/go.mod h1:HhWAlUXvxKThfOlZbcuFzsqwtF5TcqS9ru3y5GfjWAc=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
@@ -72,8 +75,13 @@ github.com/redis/go-redis/v9 v9.7.0 h1:HhLSs+B6O021gwzl+locl0zEDnyNkxMtf/Z3NNBMa
 github.com/redis/go-redis/v9 v9.7.0/go.mod h1:f6zhXITC7JUJIlPEiBOTXxJgPLdZcA93GewI7inzyWw=
 github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
 github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys=
 github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs=
+github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
+github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

go.work.sum

@@ -81,7 +81,9 @@ github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1
 github.com/rjeczalik/notify v0.9.2/go.mod h1:aErll2f0sUX9PXZnVNyeiObbmTlk5jnMoCa4QEjJeqM=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
+github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
 github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/supranational/blst v0.3.12/go.mod h1:jZJtfjgudtNl4en1tzwPIV3KjUnQUvG3/j+w+fVonLw=
 github.com/supranational/blst v0.3.13/go.mod h1:jZJtfjgudtNl4en1tzwPIV3KjUnQUvG3/j+w+fVonLw=

handler.go

@@ -78,7 +78,7 @@ type handler struct {
 	cache        Cache
 	store        Store
 	limitCounter httprate.LimitCounter
-	keyVersion   int
+	keyVersion   byte
 }
 
 var _ proto.QuotaControl = &handler{}

proto/internal/encoding/encoding.go

@@ -4,17 +4,19 @@ import (
 	"crypto/rand"
 	"encoding/binary"
 	"fmt"
+	"math"
 
 	"github.com/goware/base64"
 	"github.com/jxskiss/base62"
 )
 
 var (
 	ErrInvalidKeyLength = fmt.Errorf("invalid access key length")
+	ErrVersionMismatch  = fmt.Errorf("version mismatch")
 )
 
 type Encoding interface {
-	Version() int
+	Version() byte
 	Encode(projectID uint64, ecosystemID uint64) string
 	Decode(accessKey string) (projectID uint64, ecosystemID uint64, err error)
 }
@@ -28,7 +30,7 @@ const (
 // V0 is the v0 encoding format for project access keys
 type V0 struct{}
 
-func (V0) Version() int { return 0 }
+func (V0) Version() byte { return 0 }
 
 func (V0) Encode(projectID uint64, ecosystemID uint64) string {
 	buf := make([]byte, sizeV0)
@@ -50,47 +52,105 @@ func (V0) Decode(accessKey string) (projectID uint64, ecosystemID uint64, err er
 
 type V1 struct{}
 
-func (V1) Version() int { return 1 }
+func (V1) Version() byte { return 1 }
 
-func (V1) Encode(projectID uint64, ecosystemID uint64) string {
+func (v V1) Encode(projectID uint64, ecosystemID uint64) string {
 	buf := make([]byte, sizeV1)
-	buf[0] = byte(1)
+	buf[0] = v.Version()
 	binary.BigEndian.PutUint64(buf[1:], projectID)
 	rand.Read(buf[9:])
 	return base64.Base64UrlEncode(buf)
 }
 
-func (V1) Decode(accessKey string) (projectID uint64, ecosystemID uint64, err error) {
+func (v V1) Decode(accessKey string) (projectID uint64, ecosystemID uint64, err error) {
 	buf, err := base64.Base64UrlDecode(accessKey)
 	if err != nil {
 		return 0, 0, fmt.Errorf("base64 decode: %w", err)
 	}
 	if len(buf) != sizeV1 {
 		return 0, 0, ErrInvalidKeyLength
 	}
+	if buf[0] != v.Version() {
+		return 0, 0, ErrVersionMismatch
+	}
 	return binary.BigEndian.Uint64(buf[1:9]), 0, nil
 }
 
 type V2 struct{}
 
-func (V2) Version() int { return 2 }
+func (V2) Version() byte { return 2 }
 
-func (V2) Encode(projectID uint64, ecosystemID uint64) string {
+func (v V2) Encode(projectID uint64, ecosystemID uint64) string {
 	buf := make([]byte, sizeV2)
-	buf[0] = byte(2)
-	binary.BigEndian.PutUint64(buf[1:], projectID)
-	binary.BigEndian.PutUint64(buf[9:], ecosystemID)
-	rand.Read(buf[17:])
+	buf[0] = v.Version()
+
+	encodedProjectID := encodeUint64(projectID)
+	encodedEcosystemID := encodeUint64(ecosystemID)
+	buf[1] = byte(len(encodedProjectID)) + (byte(len(encodedEcosystemID) << 4))
+	copy(buf[2:], encodedProjectID)
+	copy(buf[2+len(encodedProjectID):], encodedEcosystemID)
+
+	rand.Read(buf[2+len(encodedProjectID)+len(encodedEcosystemID):])
+
 	return base64.Base64UrlEncode(buf)
 }
 
-func (V2) Decode(accessKey string) (projectID uint64, ecosystemID uint64, err error) {
+func (v V2) Decode(accessKey string) (projectID uint64, ecosystemID uint64, err error) {
 	buf, err := base64.Base64UrlDecode(accessKey)
 	if err != nil {
 		return 0, 0, fmt.Errorf("base64 decode: %w", err)
 	}
 	if len(buf) != sizeV2 {
 		return 0, 0, ErrInvalidKeyLength
 	}
-	return binary.BigEndian.Uint64(buf[1:9]), binary.BigEndian.Uint64(buf[9:17]), nil
+	if buf[0] != v.Version() {
+		return 0, 0, fmt.Errorf("version mismatch")
+	}
+
+	projectLength := buf[1] & 0x0f
+	ecosystemLength := buf[1] >> 4
+
+	if projectID, err = decodeUint64(buf[2 : 2+projectLength]); err != nil {
+		return 0, 0, fmt.Errorf("decode projectID: %w", err)
+	}
+
+	if ecosystemID, err = decodeUint64(buf[2+projectLength : 2+projectLength+ecosystemLength]); err != nil {
+		return 0, 0, fmt.Errorf("decode ecosystemID: %w", err)
+	}
+
+	return projectID, ecosystemID, nil
+}
+
+func encodeUint64(n uint64) []byte {
+	switch {
+	case n <= math.MaxUint8:
+		return []byte{byte(n)}
+	case n <= math.MaxUint16:
+		buf := make([]byte, 2)
+		binary.BigEndian.PutUint16(buf, uint16(n))
+		return buf
+	case n <= math.MaxUint32:
+		buf := make([]byte, 4)
+		binary.BigEndian.PutUint32(buf, uint32(n))
+		return buf
+	default:
+		buf := make([]byte, 8)
+		binary.BigEndian.PutUint64(buf, uint64(n))
+		return buf
+	}
+}
+
+func decodeUint64(buf []byte) (uint64, error) {
+	switch len(buf) {
+	case 1:
+		return uint64(buf[0]), nil
+	case 2:
+		return uint64(binary.BigEndian.Uint16(buf)), nil
+	case 4:
+		return uint64(binary.BigEndian.Uint32(buf)), nil
+	case 8:
+		return uint64(binary.BigEndian.Uint64(buf)), nil
+	default:
+		return 0, fmt.Errorf("invalid uint64 length")
+	}
 }

proto/proto.go

@@ -15,7 +15,7 @@ func Ptr[T any](v T) *T {
 	return &v
 }
 
-var supportedEncodings = []encoding.Encoding{
+var SupportedEncodings = []encoding.Encoding{
 	encoding.V0{},
 	encoding.V1{},
 	encoding.V2{},
@@ -25,19 +25,22 @@ var AccessKeyVersion = encoding.V2{}.Version()
 
 func GetProjectID(accessKey string) (projectID, ecosystemID uint64, err error) {
 	var errs []error
-	for _, e := range supportedEncodings {
+	for i := len(SupportedEncodings) - 1; i >= 0; i-- {
+		e := SupportedEncodings[i]
+
 		projectID, ecosystemID, err := e.Decode(accessKey)
 		if err != nil {
 			errs = append(errs, fmt.Errorf("decode v%d: %w", e.Version(), err))
 			continue
 		}
+
 		return projectID, ecosystemID, nil
 	}
 	return 0, 0, errors.Join(errs...)
 }
 
-func GenerateAccessKey(version int, projectID, ecosystemID uint64) string {
-	for _, e := range supportedEncodings {
+func GenerateAccessKey(version byte, projectID, ecosystemID uint64) string {
+	for _, e := range SupportedEncodings {
 		if e.Version() == version {
 			return e.Encode(projectID, ecosystemID)
 		}

proto/proto_test.go

@@ -1,6 +1,8 @@
 package proto_test
 
 import (
+	"fmt"
+	"math"
 	"testing"
 
 	"github.com/0xsequence/quotacontrol/proto"
@@ -10,37 +12,34 @@ import (
 )
 
 func TestAccessKeyEncoding(t *testing.T) {
-	t.Run("v0", func(t *testing.T) {
-		projectID := uint64(12345)
-		accessKey := proto.GenerateAccessKey(0, projectID, 0)
-		t.Log("=> k", accessKey)
-
-		outID, outecosystemID, err := proto.GetProjectID(accessKey)
-		require.NoError(t, err)
-		require.Equal(t, projectID, outID)
-		require.Equal(t, uint64(0), outecosystemID)
-	})
-
-	t.Run("v1", func(t *testing.T) {
-		projectID := uint64(12345)
-		accessKey := proto.GenerateAccessKey(1, projectID, 0)
-		t.Log("=> k", accessKey)
-		outID, ecosystemID, err := proto.GetProjectID(accessKey)
-		require.NoError(t, err)
-		require.Equal(t, projectID, outID)
-		require.Equal(t, uint64(0), ecosystemID)
-	})
-	t.Run("v1", func(t *testing.T) {
-		projectID := uint64(12345)
-		ecosystemID := uint64(54321)
-		accessKey := proto.GenerateAccessKey(2, projectID, ecosystemID)
-		t.Log("=> k", accessKey)
+	inputList := [][]uint64{
+		{1, 2},
+		{127, 128},
+		{math.MaxUint8, math.MaxUint8 + 1},
+		{math.MaxUint16, math.MaxUint16 + 1},
+		{math.MaxUint32, math.MaxUint32 + 1},
+		{math.MaxUint64, 1},
+	}
+	for _, e := range proto.SupportedEncodings {
+		version := e.Version()
+		t.Run(fmt.Sprintf("v%d", version), func(t *testing.T) {
+			for _, input := range inputList {
+				projectID := input[0]
+				ecosystemID := input[1]
+				accessKey := proto.GenerateAccessKey(version, projectID, ecosystemID)
+				t.Logf("=> key: [%d/%d] %s", projectID, ecosystemID, accessKey)
 
-		outID, outecosystemID, err := proto.GetProjectID(accessKey)
-		require.NoError(t, err)
-		require.Equal(t, projectID, outID)
-		require.Equal(t, ecosystemID, outecosystemID)
-	})
+				outID, outecosystemID, err := proto.GetProjectID(accessKey)
+				require.NoError(t, err)
+				require.Equal(t, projectID, outID)
+				if version < 2 {
+					require.Equal(t, uint64(0), outecosystemID)
+				} else {
+					require.Equal(t, ecosystemID, outecosystemID)
+				}
+			}
+		})
+	}
 }
 
 func TestAccessKeyValidateOrigin(t *testing.T) {