frostfs-http-gw/tokens/bearer-token_test.go

//go:build !integration

package tokens

import (
	"context"
	"encoding/base64"
	"testing"

	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/bearer"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/user"
	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
	"github.com/stretchr/testify/require"
	"github.com/valyala/fasthttp"
)

func makeTestCookie(value []byte) *fasthttp.RequestHeader {
	header := new(fasthttp.RequestHeader)
	header.SetCookie(bearerTokenHdr, string(value))
	return header
}

func makeTestHeader(value []byte) *fasthttp.RequestHeader {
	header := new(fasthttp.RequestHeader)
	if value != nil {
		header.Set(fasthttp.HeaderAuthorization, string(value))
	}
	return header
}

func makeBearer(value string) string {
	return bearerTokenHdr + " " + value
}

func TestBearerTokenFromCookie(t *testing.T) {
	cases := []struct {
		name   string
		actual []byte
		expect []byte
	}{
		{
			name: "empty",
		},
		{
			name:   "normal",
			actual: []byte("TOKEN"),
			expect: []byte("TOKEN"),
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.expect, BearerTokenFromCookie(makeTestCookie(tt.actual)))
		})
	}
}

func TestBearerTokenFromHeader(t *testing.T) {
	validToken := "token"
	tokenWithoutPrefix := "invalid-token"

	cases := []struct {
		name   string
		actual []byte
		expect []byte
	}{
		{
			name: "empty",
		},
		{
			name:   "token without the bearer prefix",
			actual: []byte(tokenWithoutPrefix),
		},
		{
			name:   "token without payload",
			actual: []byte(makeBearer("")),
		},
		{
			name:   "normal",
			actual: []byte(makeBearer(validToken)),
			expect: []byte(validToken),
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.expect, BearerTokenFromHeader(makeTestHeader(tt.actual)))
		})
	}
}

func TestFetchBearerToken(t *testing.T) {
	key, err := keys.NewPrivateKey()
	require.NoError(t, err)
	var uid user.ID
	user.IDFromKey(&uid, key.PrivateKey.PublicKey)

	tkn := new(bearer.Token)
	tkn.ForUser(uid)

	t64 := base64.StdEncoding.EncodeToString(tkn.Marshal())
	require.NotEmpty(t, t64)

	cases := []struct {
		name   string
		cookie string
		header string
		error  string
		nilCtx bool
		expect *bearer.Token
	}{
		{
			name: "empty",
		},
		{
			name:   "nil context",
			nilCtx: true,
		},
		{
			name:   "bad base64 header",
			header: "WRONG BASE64",
			error:  "can't base64-decode bearer token",
		},
		{
			name:   "bad base64 cookie",
			cookie: "WRONG BASE64",
			error:  "can't base64-decode bearer token",
		},
		{
			name:   "header token unmarshal error",
			header: "dGVzdAo=",
			error:  "can't unmarshal bearer token",
		},
		{
			name:   "cookie token unmarshal error",
			cookie: "dGVzdAo=",
			error:  "can't unmarshal bearer token",
		},
		{
			name:   "bad header and cookie",
			header: "WRONG BASE64",
			cookie: "dGVzdAo=",
			error:  "can't unmarshal bearer token",
		},
		{
			name:   "bad header, but good cookie",
			header: "dGVzdAo=",
			cookie: t64,
			expect: tkn,
		},
		{
			name:   "bad cookie, but good header",
			header: t64,
			cookie: "dGVzdAo=",
			expect: tkn,
		},
		{
			name:   "ok for header",
			header: t64,
			expect: tkn,
		},
		{
			name:   "ok for cookie",
			cookie: t64,
			expect: tkn,
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			var ctx *fasthttp.RequestCtx
			if !tt.nilCtx {
				ctx = makeTestRequest(tt.cookie, tt.header)
			}

			actual, err := fetchBearerToken(ctx)

			if tt.error == "" {
				require.NoError(t, err)
				require.Equal(t, tt.expect, actual)

				return
			}

			require.Contains(t, err.Error(), tt.error)
		})
	}
}

func makeTestRequest(cookie, header string) *fasthttp.RequestCtx {
	ctx := new(fasthttp.RequestCtx)

	if cookie != "" {
		ctx.Request.Header.SetCookie(bearerTokenHdr, cookie)
	}

	if header != "" {
		ctx.Request.Header.Set(fasthttp.HeaderAuthorization, bearerTokenHdr+" "+header)
	}
	return ctx
}

func TestCheckAndPropagateBearerToken(t *testing.T) {
	key, err := keys.NewPrivateKey()
	require.NoError(t, err)
	var uid user.ID
	user.IDFromKey(&uid, key.PrivateKey.PublicKey)

	tkn := new(bearer.Token)
	tkn.ForUser(uid)

	t64 := base64.StdEncoding.EncodeToString(tkn.Marshal())
	require.NotEmpty(t, t64)

	req := makeTestRequest(t64, "")

	// Expect to see the token within the context.
	appCtx, err := StoreBearerTokenAppCtx(context.Background(), req)
	require.NoError(t, err)

	// Expect to see the same token without errors.
	actual, err := LoadBearerToken(appCtx)
	require.NoError(t, err)
	require.Equal(t, tkn, actual)
}

func TestLoadBearerToken(t *testing.T) {
	ctx := context.Background()
	token := new(bearer.Token)

	cases := []struct {
		name   string
		appCtx context.Context
		error  string
	}{
		{
			name:   "token is missing in the context",
			appCtx: ctx,
			error:  "found empty bearer token",
		},
		{
			name:   "normal",
			appCtx: context.WithValue(ctx, bearerTokenKey, token),
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			tkn, err := LoadBearerToken(tt.appCtx)

			if tt.error == "" {
				require.NoError(t, err)
				require.Equal(t, token, tkn)

				return
			}

			require.Contains(t, err.Error(), tt.error)
		})
	}
}

func TestStoreBearerTokenAppCtx(t *testing.T) {
	key, err := keys.NewPrivateKey()
	require.NoError(t, err)
	var uid user.ID
	user.IDFromKey(&uid, key.PrivateKey.PublicKey)

	tkn := new(bearer.Token)
	tkn.ForUser(uid)

	t64 := base64.StdEncoding.EncodeToString(tkn.Marshal())
	require.NotEmpty(t, t64)

	cases := []struct {
		name  string
		req   *fasthttp.RequestCtx
		error string
	}{
		{
			name:  "invalid token",
			req:   makeTestRequest("dGVzdAo=", ""),
			error: "can't unmarshal bearer token",
		},
		{
			name: "normal",
			req:  makeTestRequest(t64, ""),
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			ctx, err := StoreBearerTokenAppCtx(context.Background(), tt.req)

			if tt.error == "" {
				require.NoError(t, err)
				actualToken, ok := ctx.Value(bearerTokenKey).(*bearer.Token)
				require.True(t, ok)
				require.Equal(t, tkn, actualToken)

				return
			}

			require.Contains(t, err.Error(), tt.error)
		})
	}
}