Add GetPrefix method to AccessKey and update tests for consistency

Author: klaidliadon

access_key.go

@@ -4,13 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"net/http"
 	"strings"
 
 	"crypto/rand"
 	"encoding/binary"
 
-	"github.com/go-chi/transport"
 	"github.com/goware/base64"
 	"github.com/jxskiss/base62"
 )
@@ -43,6 +41,14 @@ func (a AccessKey) GetProjectID() (projectID uint64, err error) {
 	return 0, errors.Join(errs...)
 }
 
+func (a AccessKey) GetPrefix() string {
+	parts := strings.Split(a.String(), Separator)
+	if len(parts) < 2 {
+		return ""
+	}
+	return strings.Join(parts[:len(parts)-1], Separator)
+}
+
 func NewAccessKey(ctx context.Context, projectID uint64) AccessKey {
 	version, ok := GetVersion(ctx)
 	if !ok {
@@ -57,26 +63,6 @@ func NewAccessKey(ctx context.Context, projectID uint64) AccessKey {
 	return ""
 }
 
-func GetAccessKeyPrefix(accessKey AccessKey) string {
-	parts := strings.Split(accessKey.String(), Separator)
-	if len(parts) < 2 {
-		return ""
-	}
-	return strings.Join(parts[:len(parts)-1], Separator)
-}
-
-func ForwardAccessKeyTransport(next http.RoundTripper) http.RoundTripper {
-	return transport.RoundTripFunc(func(req *http.Request) (resp *http.Response, err error) {
-		r := transport.CloneRequest(req)
-
-		if accessKey, ok := GetAccessKey(req.Context()); ok {
-			r.Header.Set(HeaderAccessKey, accessKey)
-		}
-
-		return next.RoundTrip(r)
-	})
-}
-
 type Encoding interface {
 	Version() byte
 	Encode(ctx context.Context, projectID uint64) AccessKey

access_key_test.go

@@ -35,15 +35,15 @@ func TestAccessKeyEncoding(t *testing.T) {
 		ctx := authcontrol.WithVersion(context.Background(), 2)
 		projectID := uint64(12345)
 		accessKey := authcontrol.NewAccessKey(ctx, projectID)
-		t.Log("=> k", accessKey, "| prefix =>", authcontrol.GetAccessKeyPrefix(accessKey))
+		t.Log("=> k", accessKey, "| prefix =>", accessKey.GetPrefix())
 		outID, err := accessKey.GetProjectID()
 		require.NoError(t, err)
 		require.Equal(t, projectID, outID)
 
 		ctx = authcontrol.WithPrefix(ctx, "newprefix:dev")
 
 		accessKey2 := authcontrol.NewAccessKey(ctx, projectID)
-		t.Log("=> k", accessKey2, "| prefix =>", authcontrol.GetAccessKeyPrefix(accessKey2))
+		t.Log("=> k", accessKey2, "| prefix =>", accessKey2.GetPrefix())
 		outID, err = accessKey2.GetProjectID()
 		require.NoError(t, err)
 		require.Equal(t, projectID, outID)
@@ -57,7 +57,7 @@ func TestAccessKeyEncoding(t *testing.T) {
 func TestDecode(t *testing.T) {
 	ctx := authcontrol.WithVersion(context.Background(), 2)
 	accessKey := authcontrol.NewAccessKey(ctx, 237)
-	t.Log("=> k", accessKey, "| prefix =>", authcontrol.GetAccessKeyPrefix(accessKey))
+	t.Log("=> k", accessKey, "| prefix =>", accessKey.GetPrefix())
 }
 
 func TestForwardAccessKeyTransport(t *testing.T) {

common.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/0xsequence/authcontrol/proto"
 	"github.com/go-chi/jwtauth/v5"
+	"github.com/go-chi/transport"
 	"github.com/lestrrat-go/jwx/v2/jwa"
 	"github.com/lestrrat-go/jwx/v2/jwt"
 )
@@ -198,3 +199,16 @@ func findProjectClaim(r *http.Request) (uint64, error) {
 		return 0, fmt.Errorf("invalid type: %T", val)
 	}
 }
+
+// ForwardAccessKeyTransport is a RoundTripper that forwards the access key from the request context to the request header.
+func ForwardAccessKeyTransport(next http.RoundTripper) http.RoundTripper {
+	return transport.RoundTripFunc(func(req *http.Request) (resp *http.Response, err error) {
+		r := transport.CloneRequest(req)
+
+		if accessKey, ok := GetAccessKey(req.Context()); ok {
+			r.Header.Set(HeaderAccessKey, accessKey.String())
+		}
+
+		return next.RoundTrip(r)
+	})
+}

middleware_test.go

@@ -325,8 +325,7 @@ func TestCustomErrHandler(t *testing.T) {
 
 	r.Handle("/*", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
 
-	var claims map[string]any
-	claims = map[string]any{"service": "client_service"}
+	claims := map[string]any{"service": "client_service"}
 
 	// Valid Request
 	ok, err := executeRequest(t, ctx, r, fmt.Sprintf("/rpc/%s/%s", ServiceName, MethodName), accessKey(AccessKey), jwt(authcontrol.S2SToken(JWTSecret, claims)))