Repository: jkroepke/openvpn-auth-oauth2
Pull request #238 (2024-03-31 09:04:31)

nit: add context to storage class
Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>#### What this PR does / why we need it

#### Which issue this PR fixes

*(optional, in `fixes #<issue number>(, fixes #<issue_number>, ...)` format, will close that issue when PR gets merged)*: fixes #

- fixes #

#### Special notes for your reviewer

#### Particularly user-facing changes

#### Checklist

Complete these before marking the PR as `ready to review`:

<!-- [Place an '[x]' (no spaces) in all applicable fields.] -->

- [ ] [DCO](https://github.com/prometheus-community/helm-charts/blob/main/CONTRIBUTING.md#sign-off-your-work) signed
- [ ] The PR title has a summary of the changes
- [ ] The PR body has a summary to reflect any significant (and particularly user-facing) changes introduced by this PR



Patch:
diff --git a/cmd/daemon/root.go b/cmd/daemon/root.go
index 64ad0ee..49b8d36 100644
--- a/cmd/daemon/root.go
+++ b/cmd/daemon/root.go
@@ -70,7 +70,7 @@ func Execute(args []string, logWriter io.Writer, version, commit, date string) i
 
 	httpClient := &http.Client{Transport: utils.NewUserAgentTransport(nil)}
 
-	storageClient := storage.New(conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires)
+	storageClient := storage.New(ctx, conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires)
 	oauth2Client := oauth2.New(logger, conf, storageClient, httpClient)
 	openvpnClient := openvpn.New(ctx, logger, conf, oauth2Client)
 
diff --git a/cmd/state/root.go b/cmd/state/root.go
index 1f2219b..97c8c55 100644
--- a/cmd/state/root.go
+++ b/cmd/state/root.go
@@ -14,8 +14,8 @@ func Execute(args []string, logWriter io.Writer, _, _, _ string) int {
 		return 1
 	}
 
-	session := state.NewEncoded(args[4])
-	if err := session.Decode(args[3]); err != nil {
+	session, err := state.NewWithEncodedToken(args[4], args[3])
+	if err != nil {
 		fmt.Fprintln(logWriter, err)
 
 		return 1
diff --git a/internal/oauth2/handler.go b/internal/oauth2/handler.go
index b1f7137..2664a68 100644
--- a/internal/oauth2/handler.go
+++ b/internal/oauth2/handler.go
@@ -62,8 +62,8 @@ func (p *Provider) oauth2Start() http.Handler {
 			return
 		}
 
-		session := state.NewEncoded(sessionState)
-		if err := session.Decode(p.conf.HTTP.Secret.String()); err != nil {
+		session, err := state.NewWithEncodedToken(sessionState, p.conf.HTTP.Secret.String())
+		if err != nil {
 			p.logger.Warn(utils.StringConcat("invalid state: ", err.Error()))
 			w.WriteHeader(http.StatusBadRequest)
 
@@ -155,8 +155,8 @@ func (p *Provider) oauth2Callback() http.Handler {
 			return
 		}
 
-		session := state.NewEncoded(encryptedState)
-		if err := session.Decode(p.conf.HTTP.Secret.String()); err != nil {
+		session, err := state.NewWithEncodedToken(encryptedState, p.conf.HTTP.Secret.String())
+		if err != nil {
 			writeError(w, p.logger, p.conf, http.StatusBadRequest, "Invalid State", err.Error())
 
 			return
diff --git a/internal/oauth2/provider.go b/internal/oauth2/provider.go
index c757b2d..642a496 100644
--- a/internal/oauth2/provider.go
+++ b/internal/oauth2/provider.go
@@ -225,8 +225,8 @@ func errorHandler(
 	w http.ResponseWriter, conf config.Config, logger *slog.Logger, openvpn OpenVPN,
 	httpStatus int, errorType string, errorDesc string, encryptedSession string,
 ) {
-	session := state.NewEncoded(encryptedSession)
-	if err := session.Decode(conf.HTTP.Secret.String()); err == nil {
+	session, err := state.NewWithEncodedToken(encryptedSession, conf.HTTP.Secret.String())
+	if err == nil {
 		logger = logger.With(
 			slog.String("ip", fmt.Sprintf("%s:%s", session.IPAddr, session.IPPort)),
 			slog.Uint64("cid", session.Client.CID),
diff --git a/internal/openvpn/client.go b/internal/openvpn/client.go
index 498f3cc..f9fbdec 100644
--- a/internal/openvpn/client.go
+++ b/internal/openvpn/client.go
@@ -90,11 +90,13 @@ func (c *Client) handleClientAuthentication(logger *slog.Logger, client connecti
 	}
 
 	session := state.New(ClientIdentifier, ipAddr, ipPort, commonName)
-	if err := session.Encode(c.conf.HTTP.Secret.String()); err != nil {
+
+	encodedSession, err := session.Encode(c.conf.HTTP.Secret.String())
+	if err != nil {
 		return fmt.Errorf("error encoding state: %w", err)
 	}
 
-	startURL := utils.StringConcat(strings.TrimSuffix(c.conf.HTTP.BaseURL.String(), "/"), "/oauth2/start?state=", session.Encoded())
+	startURL := utils.StringConcat(strings.TrimSuffix(c.conf.HTTP.BaseURL.String(), "/"), "/oauth2/start?state=", encodedSession)
 
 	if len(startURL) >= 245 {
 		c.DenyClient(logger, ClientIdentifier, "internal error")
@@ -106,7 +108,7 @@ func (c *Client) handleClientAuthentication(logger *slog.Logger, client connecti
 
 	logger.Info("start pending auth")
 
-	_, err := c.SendCommandf(`client-pending-auth %d %d "WEB_AUTH::%s" %.0f`, client.CID, client.KID, startURL, c.conf.OpenVpn.AuthPendingTimeout.Seconds())
+	_, err = c.SendCommandf(`client-pending-auth %d %d "WEB_AUTH::%s" %.0f`, client.CID, client.KID, startURL, c.conf.OpenVpn.AuthPendingTimeout.Seconds())
 	if err != nil {
 		logger.Warn(err.Error())
 	}
@@ -139,12 +141,7 @@ func (c *Client) checkReauth(logger *slog.Logger, client connection.Client) bool
 		return true
 	}
 
-	var (
-		err error
-		ok  bool
-	)
-
-	ok, err = c.oauth2.RefreshClientAuth(logger, client)
+	ok, err := c.oauth2.RefreshClientAuth(logger, client)
 	if err != nil {
 		logger.Warn(err.Error())
 	}
diff --git a/internal/state/state.go b/internal/state/state.go
index 9f1d7a4..3afb637 100644
--- a/internal/state/state.go
+++ b/internal/state/state.go
@@ -17,8 +17,6 @@ type State struct {
 	IPPort     string
 	CommonName string
 	Issued     int64
-
-	encoded string
 }
 
 type ClientIdentifier struct {
@@ -39,25 +37,25 @@ func New(client ClientIdentifier, ipAddr, ipPort, commonName string) State {
 	}
 }
 
-func NewEncoded(state string) State {
-	return State{
-		encoded: state,
+func NewWithEncodedToken(encodedState, secretKey string) (State, error) {
+	state := State{}
+
+	if err := state.decode(encodedState, secretKey); err != nil {
+		return State{}, err
 	}
-}
 
-func (state *State) Encoded() string {
-	return state.encoded
+	return state, nil
 }
 
-func (state *State) Decode(secretKey string) error {
-	encrypted, err := base64.RawURLEncoding.DecodeString(state.encoded)
+func (state *State) decode(encodedState, secretKey string) error {
+	encrypted, err := base64.RawURLEncoding.DecodeString(encodedState)
 	if err != nil {
-		return fmt.Errorf("base64 decode %s: %w", state.encoded, err)
+		return fmt.Errorf("base64 decode %s: %w", encodedState, err)
 	}
 
 	data, err := crypto.DecryptBytesAES(encrypted, secretKey)
 	if err != nil {
-		return fmt.Errorf("decrypt aes %s: %w", state.encoded, err)
+		return fmt.Errorf("decrypt aes %s: %w", encodedState, err)
 	}
 
 	_, err = fmt.Fscanln(bytes.NewReader(data),
@@ -93,7 +91,7 @@ func (state *State) Decode(secretKey string) error {
 	return nil
 }
 
-func (state *State) Encode(secretKey string) error {
+func (state *State) Encode(secretKey string) (string, error) {
 	var data bytes.Buffer
 
 	data.Grow(512)
@@ -118,12 +116,10 @@ func (state *State) Encode(secretKey string) error {
 
 	encrypted, err := crypto.EncryptBytesAES(data.Bytes(), secretKey)
 	if err != nil {
-		return fmt.Errorf("encrypt aes: %w", err)
+		return "", fmt.Errorf("encrypt aes: %w", err)
 	}
 
-	state.encoded = base64.RawURLEncoding.EncodeToString(encrypted)
-
-	return nil
+	return base64.RawURLEncoding.EncodeToString(encrypted), nil
 }
 
 func encodeString(text string) string {
diff --git a/internal/storage/storage.go b/internal/storage/storage.go
index 4ef0ffa..0ce67a9 100644
--- a/internal/storage/storage.go
+++ b/internal/storage/storage.go
@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -20,33 +21,36 @@ type item struct {
 	expires time.Time
 }
 
-func New(encryptionKey string, expires time.Duration) *Storage {
+func New(ctx context.Context, encryptionKey string, expires time.Duration) *Storage {
 	storage := &Storage{
 		encryptionKey,
 		expires,
 		sync.Map{},
 	}
-	go storage.collect()
+	go storage.collect(ctx)
 
 	return storage
 }
 
-func (s *Storage) collect() {
+func (s *Storage) collect(ctx context.Context) {
 	for {
-		time.Sleep(time.Minute * 5)
-
-		s.data.Range(func(client, data any) bool {
-			entry, ok := data.(item)
-			if !ok {
-				panic(data)
-			}
-
-			if entry.expires.Compare(time.Now()) == -1 {
-				s.data.Delete(client)
-			}
-
-			return true
-		})
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(time.Minute * 5):
+			s.data.Range(func(client, data any) bool {
+				entry, ok := data.(item)
+				if !ok {
+					panic(data)
+				}
+
+				if entry.expires.Compare(time.Now()) == -1 {
+					s.data.Delete(client)
+				}
+
+				return true
+			})
+		}
 	}
 }
 
diff --git a/pkg/plugin/plugin.go b/pkg/plugin/plugin.go
index 471f143..da7ce66 100644
--- a/pkg/plugin/plugin.go
+++ b/pkg/plugin/plugin.go
@@ -77,10 +77,11 @@ func openvpn_plugin_open_v3_go(v3structver C.int, args *C.struct_openvpn_plugin_
 		conf:   conf,
 	}
 
-	storageClient := storage.New(conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires)
+	ctx := context.Background()
+	storageClient := storage.New(ctx, conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires)
 	provider := oauth2.New(logger, conf, storageClient, http2.DefaultClient)
 
-	if err = provider.Initialize(context.Background(), handle); err != nil {
+	if err = provider.Initialize(ctx, handle); err != nil {
 		logger.Error(err.Error())
 
 		return C.OPENVPN_PLUGIN_FUNC_ERROR
@@ -129,14 +130,15 @@ func openvpn_plugin_func_v3_go(v3structver C.int, args *C.struct_openvpn_plugin_
 		AuthFailedReasonFile: client.AuthFailedReasonFile,
 	}
 	session := state.New(clientIdentifier, client.IpAddr, client.IpPort, client.CommonName)
-	if err := session.Encode(handle.conf.HTTP.Secret.String()); err != nil {
+	encodedSession, err := session.Encode(handle.conf.HTTP.Secret.String())
+	if err != nil {
 		handle.logger.Error(fmt.Errorf("encoding state: %w", err).Error())
 		return C.OPENVPN_PLUGIN_FUNC_ERROR
 	}
 
 	startURL := utils.StringConcat(
 		strings.TrimSuffix(handle.conf.HTTP.BaseURL.String(), "/"),
-		"/oauth2/start?state=", url.QueryEscape(session.Encoded()),
+		"/oauth2/start?state=", url.QueryEscape(encodedSession),
 	)
 
 	pendingAuth := fmt.Sprintf("%.0f\nwebauth\nWEB_AUTH::%s", handle.conf.OpenVpn.AuthPendingTimeout.Seconds(), startURL)


Test patch:
diff --git a/internal/oauth2/handler_test.go b/internal/oauth2/handler_test.go
index 4d7dedc..84488da 100644
--- a/internal/oauth2/handler_test.go
+++ b/internal/oauth2/handler_test.go
@@ -384,17 +384,17 @@ func TestHandler(t *testing.T) {
 
 				time.Sleep(time.Millisecond * 100)
 
+				var err error
+
 				session := tt.state
 
 				if tt.state == "-" {
 					sessionState := state.New(state.ClientIdentifier{CID: 0, KID: 1}, tt.ipaddr, "12345", "name")
-					err := sessionState.Encode(tt.conf.HTTP.Secret.String())
+					session, err = sessionState.Encode(tt.conf.HTTP.Secret.String())
 
 					if !assert.NoError(t, err) {
 						return
 					}
-
-					session = sessionState.Encoded()
 				}
 
 				request, err := http.NewRequestWithContext(context.Background(), http.MethodGet,
diff --git a/internal/oauth2/provider_test.go b/internal/oauth2/provider_test.go
index 4dc988d..b42773c 100644
--- a/internal/oauth2/provider_test.go
+++ b/internal/oauth2/provider_test.go
@@ -142,9 +142,11 @@ func TestNewProvider(t *testing.T) {
 
 			tt.conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, tt.conf, storageClient, http2.DefaultClient)
-			ctx := context.Background()
 
 			client := openvpn.New(ctx, logger.Logger, tt.conf, provider)
 			defer client.Shutdown()
diff --git a/internal/openvpn/bench_test.go b/internal/openvpn/bench_test.go
index 3412471..31edf0a 100644
--- a/internal/openvpn/bench_test.go
+++ b/internal/openvpn/bench_test.go
@@ -45,8 +45,11 @@ func BenchmarkOpenVPNHandler(b *testing.B) {
 		},
 	}
 
-	storageClient := storage.New(testutils.Secret, time.Hour)
-	openVPNClient := openvpn.New(context.Background(), logger.Logger, conf, oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient))
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Hour)
+	openVPNClient := openvpn.New(ctx, logger.Logger, conf, oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient))
 
 	wg := sync.WaitGroup{}
 	wg.Add(1)
@@ -126,16 +129,17 @@ func BenchmarkOpenVPNPassthrough(b *testing.B) {
 
 	passThroughInterface.Close()
 
-	storageClient := storage.New(testutils.Secret, time.Hour)
+	ctx, cancel := context.WithCancelCause(context.Background())
+	defer cancel(nil)
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 	provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-	openVPNClient := openvpn.New(context.Background(), logger.Logger, conf, provider)
+	openVPNClient := openvpn.New(ctx, logger.Logger, conf, provider)
 
 	defer openVPNClient.Shutdown()
 
 	wg := sync.WaitGroup{}
 
-	ctx, cancel := context.WithCancelCause(context.Background())
-
 	wg.Add(1)
 
 	go func() {
diff --git a/internal/openvpn/main_test.go b/internal/openvpn/main_test.go
index 5c16d9d..fd2863f 100644
--- a/internal/openvpn/main_test.go
+++ b/internal/openvpn/main_test.go
@@ -39,9 +39,13 @@ func TestClientInvalidServer(t *testing.T) {
 			Bypass: config.OpenVpnBypass{CommonNames: make([]string, 0)},
 		},
 	}
-	storageClient := storage.New(testutils.Secret, time.Hour)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 	provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-	client := openvpn.New(context.Background(), logger.Logger, conf, provider)
+	client := openvpn.New(ctx, logger.Logger, conf, provider)
 	err := client.Connect()
 	require.Error(t, err)
 	assert.Equal(t, "unable to connect to openvpn management interface tcp://0.0.0.0:1: dial tcp 0.0.0.0:1: connect: connection refused", err.Error())
@@ -255,9 +259,12 @@ func TestClientFull(t *testing.T) {
 
 			tt.conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, tt.conf, storageClient, http2.DefaultClient)
-			client := openvpn.New(context.Background(), logger.Logger, tt.conf, provider)
+			client := openvpn.New(ctx, logger.Logger, tt.conf, provider)
 
 			wg := sync.WaitGroup{}
 			wg.Add(1)
@@ -304,8 +311,7 @@ func TestClientFull(t *testing.T) {
 					matches := regexp.MustCompile(`state=(.+)"`).FindStringSubmatch(auth)
 					assert.Len(t, matches, 2)
 
-					sessionState := state.NewEncoded(matches[1])
-					err := sessionState.Decode(tt.conf.HTTP.Secret.String())
+					sessionState, err := state.NewWithEncodedToken(matches[1], tt.conf.HTTP.Secret.String())
 					require.NoError(t, err) //nolint:testifylint
 
 					assert.Equal(t, uint64(1), sessionState.Client.CID)
@@ -354,9 +360,12 @@ func TestClientInvalidPassword(t *testing.T) {
 		},
 	}
 
-	storageClient := storage.New(testutils.Secret, time.Hour)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 	provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-	client := openvpn.New(context.Background(), logger.Logger, conf, provider)
+	client := openvpn.New(ctx, logger.Logger, conf, provider)
 
 	go func() {
 		conn, err := managementInterface.Accept()
@@ -428,11 +437,13 @@ func TestClientInvalidVersion(t *testing.T) {
 
 			conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			ctx, cancel := context.WithCancelCause(context.Background())
+			defer cancel(nil)
+
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-			openVPNClient := openvpn.New(context.Background(), logger.Logger, conf, provider)
+			openVPNClient := openvpn.New(ctx, logger.Logger, conf, provider)
 
-			ctx, cancel := context.WithCancelCause(context.Background())
 			wg := sync.WaitGroup{}
 			wg.Add(2)
 
@@ -507,9 +518,12 @@ func TestSIGHUP(t *testing.T) {
 
 	conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-	storageClient := storage.New(testutils.Secret, time.Hour)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 	provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-	client := openvpn.New(context.Background(), logger.Logger, conf, provider)
+	client := openvpn.New(ctx, logger.Logger, conf, provider)
 
 	wg := sync.WaitGroup{}
 	wg.Add(1)
@@ -577,9 +591,12 @@ func TestDeadLocks(t *testing.T) {
 
 			conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-			client := openvpn.New(context.Background(), logger.Logger, conf, provider)
+			client := openvpn.New(ctx, logger.Logger, conf, provider)
 
 			wg := sync.WaitGroup{}
 			wg.Add(1)
@@ -646,9 +663,12 @@ func TestInvalidCommandResponses(t *testing.T) {
 
 			conf.OpenVpn.Addr = &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, conf, storageClient, http2.DefaultClient)
-			client := openvpn.New(context.Background(), logger.Logger, conf, provider)
+			client := openvpn.New(ctx, logger.Logger, conf, provider)
 
 			wg := sync.WaitGroup{}
 			wg.Add(1)
diff --git a/internal/openvpn/passthrough_test.go b/internal/openvpn/passthrough_test.go
index abfa22a..aa7ec0f 100644
--- a/internal/openvpn/passthrough_test.go
+++ b/internal/openvpn/passthrough_test.go
@@ -182,8 +182,9 @@ func TestPassthroughFull(t *testing.T) {
 			passThroughInterface.Close()
 
 			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
 
-			storageClient := storage.New(testutils.Secret, time.Hour)
+			storageClient := storage.New(ctx, testutils.Secret, time.Hour)
 			provider := oauth2.New(logger.Logger, tt.conf, storageClient, http.DefaultClient)
 			openVPNClient := openvpn.New(ctx, logger.Logger, tt.conf, provider)
 
diff --git a/internal/state/bench_test.go b/internal/state/bench_test.go
index fa3e757..4c9fe88 100644
--- a/internal/state/bench_test.go
+++ b/internal/state/bench_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/jkroepke/openvpn-auth-oauth2/internal/state"
 	"github.com/jkroepke/openvpn-auth-oauth2/pkg/testutils"
+	"github.com/stretchr/testify/require"
 )
 
 func BenchmarkState(b *testing.B) {
@@ -17,7 +18,7 @@ func BenchmarkState(b *testing.B) {
 
 	b.Run("encode", func(b *testing.B) {
 		for range b.N {
-			_ = token.Encode(encryptionKey)
+			_, _ = token.Encode(encryptionKey)
 		}
 
 		b.ReportAllocs()
@@ -25,13 +26,14 @@ func BenchmarkState(b *testing.B) {
 
 	b.StopTimer()
 
-	encodedToken := state.NewEncoded(token.Encoded())
+	encodedTokenString, err := token.Encode(encryptionKey)
+	require.NoError(b, err)
 
 	b.StartTimer()
 
 	b.Run("decode", func(b *testing.B) {
 		for range b.N {
-			_ = encodedToken.Decode(encryptionKey)
+			_, _ = state.NewWithEncodedToken(encodedTokenString, encryptionKey)
 		}
 
 		b.ReportAllocs()
diff --git a/internal/state/state_test.go b/internal/state/state_test.go
index b6dbb41..70b6729 100644
--- a/internal/state/state_test.go
+++ b/internal/state/state_test.go
@@ -17,10 +17,11 @@ func TestState(t *testing.T) {
 
 	for i := 1; i < 50; i++ {
 		token := state.New(state.ClientIdentifier{CID: 9223372036854775807, KID: 2}, "127.0.0.1", "12345", "test")
-		require.NoError(t, token.Encode(encryptionKey))
+		encodedTokenString, err := token.Encode(encryptionKey)
+		require.NoError(t, err)
 
-		encodedToken := state.NewEncoded(token.Encoded())
-		require.NoError(t, encodedToken.Decode(encryptionKey))
+		encodedToken, err := state.NewWithEncodedToken(encodedTokenString, encryptionKey)
+		require.NoError(t, err)
 
 		assert.Equal(t, token.Client.CID, encodedToken.Client.CID)
 		assert.Equal(t, token.Client.KID, encodedToken.Client.KID)
@@ -35,10 +36,11 @@ func TestStateWithEmptyValues(t *testing.T) {
 	encryptionKey := testutils.Secret
 
 	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "")
-	require.NoError(t, token.Encode(encryptionKey))
+	encodedTokenString, err := token.Encode(encryptionKey)
+	require.NoError(t, err)
 
-	encodedToken := state.NewEncoded(token.Encoded())
-	require.NoError(t, encodedToken.Decode(encryptionKey))
+	encodedToken, err := state.NewWithEncodedToken(encodedTokenString, encryptionKey)
+	require.NoError(t, err)
 
 	assert.Equal(t, token.Client.CID, encodedToken.Client.CID)
 	assert.Equal(t, token.Client.KID, encodedToken.Client.KID)
@@ -52,19 +54,9 @@ func TestStateInvalid_Key(t *testing.T) {
 	encryptionKey := "01234567891011"
 
 	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "test")
-	require.Error(t, token.Encode(encryptionKey), "crypto/aes: invalid key size 14")
-}
-
-func TestStateInvalid_Future(t *testing.T) {
-	t.Parallel()
+	_, err := token.Encode(encryptionKey)
 
-	encryptionKey := testutils.Secret
-
-	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "test")
-	token.Issued = time.Now().Add(time.Hour).Unix()
-
-	require.NoError(t, token.Encode(encryptionKey))
-	assert.Contains(t, token.Decode(encryptionKey).Error(), "invalid state: issued in future, issued at:")
+	require.Error(t, err, "crypto/aes: invalid key size 14")
 }
 
 func TestState_WithSpace(t *testing.T) {
@@ -74,10 +66,12 @@ func TestState_WithSpace(t *testing.T) {
 
 	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "te st")
 
-	require.NoError(t, token.Encode(encryptionKey))
+	encodedTokenString, err := token.Encode(encryptionKey)
+
+	require.NoError(t, err)
 
-	encodedToken := state.NewEncoded(token.Encoded())
-	require.NoError(t, encodedToken.Decode(encryptionKey))
+	encodedToken, err := state.NewWithEncodedToken(encodedTokenString, encryptionKey)
+	require.NoError(t, err)
 
 	assert.Equal(t, token.Client.CID, encodedToken.Client.CID)
 	assert.Equal(t, token.Client.KID, encodedToken.Client.KID)
@@ -85,6 +79,22 @@ func TestState_WithSpace(t *testing.T) {
 	assert.Equal(t, token.CommonName, encodedToken.CommonName)
 }
 
+func TestStateInvalid_Future(t *testing.T) {
+	t.Parallel()
+
+	encryptionKey := testutils.Secret
+
+	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "test")
+	token.Issued = time.Now().Add(time.Hour).Unix()
+	encodedTokenString, err := token.Encode(encryptionKey)
+
+	require.NoError(t, err)
+
+	_, err = state.NewWithEncodedToken(encodedTokenString, encryptionKey)
+
+	assert.ErrorContains(t, err, "invalid state: issued in future, issued at:")
+}
+
 func TestStateInvalid_TooOld(t *testing.T) {
 	t.Parallel()
 
@@ -92,16 +102,18 @@ func TestStateInvalid_TooOld(t *testing.T) {
 
 	token := state.New(state.ClientIdentifier{CID: 1, KID: 2}, "127.0.0.1", "12345", "test")
 	token.Issued = time.Now().Add(-1 * time.Hour).Unix()
+	encodedTokenString, err := token.Encode(encryptionKey)
+
+	require.NoError(t, err)
 
-	require.NoError(t, token.Encode(encryptionKey))
-	assert.Contains(t, token.Decode(encryptionKey).Error(), "invalid state: expired after 2 minutes, issued at:")
+	_, err = state.NewWithEncodedToken(encodedTokenString, encryptionKey)
+
+	assert.ErrorContains(t, err, "invalid state: expired after 2 minutes, issued at:")
 }
 
 func TestStateInvalid_Encoded(t *testing.T) {
 	t.Parallel()
 
-	encryptionKey := testutils.Secret
-
-	token := state.NewEncoded("test")
-	require.Error(t, token.Decode(encryptionKey))
+	_, err := state.NewWithEncodedToken("test", testutils.Secret)
+	require.Error(t, err)
 }
diff --git a/internal/storage/storage_test.go b/internal/storage/storage_test.go
index 5d134b9..7db46d5 100644
--- a/internal/storage/storage_test.go
+++ b/internal/storage/storage_test.go
@@ -1,6 +1,7 @@
 package storage_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -13,7 +14,10 @@ import (
 func TestStorage(t *testing.T) {
 	t.Parallel()
 
-	storageClient := storage.New(testutils.Secret, time.Millisecond*400)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	storageClient := storage.New(ctx, testutils.Secret, time.Millisecond*400)
 	require.NoError(t, storageClient.Set("0", "TEST0"))
 	require.NoError(t, storageClient.Set("1", "TEST1"))
 
diff --git a/pkg/testutils/main.go b/pkg/testutils/main.go
index ab9eda2..d22d6c2 100644
--- a/pkg/testutils/main.go
+++ b/pkg/testutils/main.go
@@ -282,7 +282,7 @@ func SetupMockEnvironment(ctx context.Context, tb testing.TB, conf config.Config
 	}
 
 	httpClient := &http.Client{Transport: NewMockRoundTripper(utils.NewUserAgentTransport(nil))}
-	storageClient := storage.New(Secret, conf.OAuth2.Refresh.Expires)
+	storageClient := storage.New(ctx, Secret, conf.OAuth2.Refresh.Expires)
 	provider := oauth2.New(logger.Logger, conf, storageClient, httpClient)
 	openvpnClient := openvpn.New(ctx, logger.Logger, conf, provider)