Repository: jkroepke/openvpn-auth-oauth2 (go)

Issue:
Add server in BYPASS without restart
### Problem Statement

Hi,

Is there a way to add/remove a server to CONFIG_OPENVPN_BYPASS_COMMON__NAMES and reload the configuration without disconnecting all other clients?

Best regards,

### openvpn-auth-oauth2 logs

```shell
None
```

### Environment

- openvpn-auth-oauth2 Version:
- OpenVPN Server Version:
- Server OS:
- OpenVPN Client (flavor, OS):


### Preflight Checklist

- [x] I could not find a solution in the [documentation](https://github.com/jkroepke/openvpn-auth-oauth2/wiki),
the [FAQ](https://github.com/jkroepke/openvpn-auth-oauth2/wiki/FAQ), the existing issues or discussions.


Pull request description:
feat: add soft-reload
#### 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 #508 

#### 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 c82e425..d22cc4c 100644
--- a/cmd/daemon/root.go
+++ b/cmd/daemon/root.go
@@ -30,10 +30,25 @@ import (
 	"github.com/jkroepke/openvpn-auth-oauth2/internal/version"
 )
 
-// Execute runs the main program logic of openvpn-auth-oauth2.
-//
-//nolint:cyclop
+var ErrReload = errors.New("reload")
+
+// Execute is the main entry point for the openvpn-auth-oauth2 daemon.
 func Execute(args []string, logWriter io.Writer) int {
+	tokenDataStorage := tokenstorage.DataMap{}
+
+	for {
+		if rt := run(args, logWriter, tokenDataStorage); rt != -1 {
+			return rt
+		}
+
+		time.Sleep(300 * time.Millisecond) // Wait before reloading configuration
+	}
+}
+
+// run runs the main program logic of openvpn-auth-oauth2.
+//
+//nolint:cyclop,gocognit
+func run(args []string, logWriter io.Writer, tokenDataStorage tokenstorage.DataMap) int {
 	conf, err := configure(args, logWriter)
 	if err != nil {
 		if errors.Is(err, flag.ErrHelp) {
@@ -64,7 +79,8 @@ func Execute(args []string, logWriter io.Writer) int {
 	logger.LogAttrs(ctx, slog.LevelDebug, "config", slog.String("config", conf.String()))
 
 	httpClient := &http.Client{Transport: utils.NewUserAgentTransport(http.DefaultTransport)}
-	tokenStorage := tokenstorage.NewInMemory(ctx, conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(conf.OAuth2.Refresh.Secret.String(), conf.OAuth2.Refresh.Expires)
+	tokenStorage.SetStorage(tokenDataStorage)
 
 	var provider oauth2.Provider
 
@@ -142,7 +158,7 @@ func Execute(args []string, logWriter io.Writer) int {
 	}()
 
 	termCh := make(chan os.Signal, 1)
-	signal.Notify(termCh, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM)
+	signal.Notify(termCh, os.Interrupt, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGUSR1)
 
 	logger.LogAttrs(ctx, slog.LevelInfo,
 		"openvpn-auth-oauth2 started with base url "+conf.HTTP.BaseURL.String(),
@@ -152,7 +168,15 @@ func Execute(args []string, logWriter io.Writer) int {
 		select {
 		case <-ctx.Done():
 			err = context.Cause(ctx)
-			if err != nil && !errors.Is(err, context.Canceled) {
+			if err != nil {
+				if errors.Is(err, context.Canceled) {
+					return 0
+				}
+
+				if errors.Is(err, ErrReload) {
+					return -1
+				}
+
 				logger.Error(err.Error())
 
 				return 1
@@ -160,13 +184,16 @@ func Execute(args []string, logWriter io.Writer) int {
 
 			return 0
 		case sig := <-termCh:
-			logger.Info("receiving signal: " + sig.String())
+			logger.LogAttrs(ctx, slog.LevelInfo, "receiving signal: "+sig.String())
 
 			switch sig {
 			case syscall.SIGHUP:
 				if err = server.Reload(); err != nil {
 					cancel(fmt.Errorf("error reloading http server: %w", err))
 				}
+			case syscall.SIGUSR1:
+				logger.LogAttrs(ctx, slog.LevelInfo, "reloading configuration")
+				cancel(ErrReload)
 			default:
 				cancel(nil)
 			}
diff --git a/internal/config/config.go b/internal/config/config.go
index 97220ae..9369872 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -6,6 +6,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"log/slog"
 	"os"
 	"strconv"
 	"strings"
@@ -196,6 +197,18 @@ func lookupEnvOrDefault[T any](key string, defaultValue T) T {
 			return defaultValue
 		}
 
+		value, ok = any(typedValue).(T)
+	case slog.Level:
+		if err := typedValue.UnmarshalText([]byte(envValue)); err != nil {
+			return defaultValue
+		}
+
+		value, ok = any(typedValue).(T)
+	case OpenVPNCommonNameMode:
+		if err := typedValue.UnmarshalText([]byte(envValue)); err != nil {
+			return defaultValue
+		}
+
 		value, ok = any(typedValue).(T)
 	case types.Template:
 		if err := typedValue.UnmarshalText([]byte(envValue)); err != nil {
diff --git a/internal/config/types.go b/internal/config/types.go
index 6ab7cf2..44a183d 100644
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -206,22 +206,21 @@ func (s OAuth2AuthStyle) MarshalText() ([]byte, error) {
 	case OAuth2AuthStyle(oauth2.AuthStyleInHeader):
 		return []byte("AuthStyleInHeader"), nil
 	default:
-		return nil, fmt.Errorf("unknown auth-style %d", s)
+		return nil, fmt.Errorf("unknown auth-style: %d", s)
 	}
 }
 
 //goland:noinspection GoMixedReceiverTypes
 func (s *OAuth2AuthStyle) UnmarshalText(text []byte) error {
-	config := strings.ToLower(string(text))
-	switch config {
-	case strings.ToLower("AuthStyleAutoDetect"):
+	switch strings.ToLower(string(text)) {
+	case "authstyleautodetect":
 		*s = OAuth2AuthStyle(oauth2.AuthStyleAutoDetect)
-	case strings.ToLower("AuthStyleInParams"):
+	case "authstyleinparams":
 		*s = OAuth2AuthStyle(oauth2.AuthStyleInParams)
-	case strings.ToLower("AuthStyleInHeader"):
+	case "authstyleinheader":
 		*s = OAuth2AuthStyle(oauth2.AuthStyleInHeader)
 	default:
-		return fmt.Errorf("unknown auth-style %d", s)
+		return fmt.Errorf("unknown auth-style: %s", text)
 	}
 
 	return nil
diff --git a/internal/oauth2/refresh.go b/internal/oauth2/refresh.go
index 52a32de..8ae94ca 100644
--- a/internal/oauth2/refresh.go
+++ b/internal/oauth2/refresh.go
@@ -116,7 +116,11 @@ func (c Client) ClientDisconnect(ctx context.Context, logger *slog.Logger, clien
 		return
 	}
 
-	c.storage.Delete(id)
+	if err = c.storage.Delete(id); err != nil {
+		logger.LogAttrs(ctx, slog.LevelWarn, fmt.Errorf("error delete refresh token from storage: %w", err).Error())
+
+		return
+	}
 
 	if !c.conf.OAuth2.Refresh.ValidateUser {
 		return
diff --git a/internal/tokenstorage/inmemory.go b/internal/tokenstorage/inmemory.go
index 19d865c..6c62fcf 100644
--- a/internal/tokenstorage/inmemory.go
+++ b/internal/tokenstorage/inmemory.go
@@ -1,7 +1,7 @@
 package tokenstorage
 
 import (
-	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -9,61 +9,88 @@ import (
 	"github.com/zitadel/oidc/v3/pkg/crypto"
 )
 
+// InMemory provides an in-memory implementation of a token storage system.
+// It stores encrypted tokens associated with clients, supports expiration, and is safe for concurrent use.
 type InMemory struct {
-	cancelFn      context.CancelFunc
-	data          sync.Map
-	encryptionKey string
-	expires       time.Duration
-}
+	data DataMap // holds the actual token data mapped by client identifier.
 
-type item struct {
-	expires time.Time
-	token   []byte
+	encryptionKey string        // used to encrypt and decrypt token data.
+	mu            sync.RWMutex  // read-write mutex to ensure safe concurrent access.
+	expires       time.Duration // defines the duration after which a token is considered expired.
 }
 
-func NewInMemory(ctx context.Context, encryptionKey string, expires, cleanupInterval time.Duration) *InMemory {
-	ctx, cancel := context.WithCancel(ctx)
-
+// NewInMemory creates a new in-memory token storage with the given encryption key and expiration duration.
+func NewInMemory(encryptionKey string, expires time.Duration) *InMemory {
 	storage := &InMemory{
-		cancelFn:      cancel,
-		data:          sync.Map{},
+		data:          DataMap{},
 		encryptionKey: encryptionKey,
 		expires:       expires,
+		mu:            sync.RWMutex{},
 	}
 
-	go storage.collect(ctx, time.NewTicker(cleanupInterval))
-
 	return storage
 }
 
+// SetStorage replaces the current storage data with the provided DataMap.
+// This is mainly used for testing or restoring state.
+func (s *InMemory) SetStorage(data DataMap) {
+	s.mu.Lock()
+
+	s.data = data
+
+	s.mu.Unlock()
+}
+
+// Set stores an encrypted token for a given client, with an expiration time.
+// The token is encrypted using AES before storage.
 func (s *InMemory) Set(client, token string) error {
+	s.mu.Lock()
+
 	encryptedBytes, err := crypto.EncryptBytesAES([]byte(token), s.encryptionKey)
 	if err != nil {
-		return fmt.Errorf("decrypt error: %w", err)
+		s.mu.Unlock()
+
+		return fmt.Errorf("encrypt error: %w", err)
+	}
+
+	s.data[client] = item{
+		Data:    encryptedBytes,
+		Expires: time.Now().Add(s.expires),
 	}
 
-	s.data.Store(client, item{token: encryptedBytes, expires: time.Now().Add(s.expires)})
+	s.mu.Unlock()
 
 	return nil
 }
 
+// Get retrieves and decrypts the token for a given client.
+// If the token is expired or does not exist, an error is returned.
 func (s *InMemory) Get(client string) (string, error) {
-	data, ok := s.data.Load(client)
-	if !ok {
+	s.mu.RLock()
+	if s.data == nil {
+		s.mu.RUnlock()
+
 		return "", ErrNotExists
 	}
 
-	item, ok := data.(item)
+	data, ok := s.data[client]
 	if !ok {
-		s.Delete(client)
+		s.mu.RUnlock()
 
 		return "", ErrNotExists
 	}
 
-	encryptedBytes := make([]byte, len(item.token))
+	s.mu.RUnlock()
+
+	if data.Expires.Before(time.Now()) {
+		delErr := s.Delete(client)
+
+		return "", errors.Join(ErrNotExists, delErr)
+	}
 
 	// we need to copy the data, since crypto.DecryptBytesAES will modify the slice in place
-	copy(encryptedBytes, item.token)
+	encryptedBytes := make([]byte, len(data.Data))
+	copy(encryptedBytes, data.Data)
 
 	token, err := crypto.DecryptBytesAES(encryptedBytes, s.encryptionKey)
 	if err != nil {
@@ -73,32 +100,25 @@ func (s *InMemory) Get(client string) (string, error) {
 	return string(token), nil
 }
 
-func (s *InMemory) Delete(client string) {
-	s.data.Delete(client)
-}
+// Delete removes the token data for a given client from storage.
+// If the storage is empty or the client does not exist, it does nothing.
+func (s *InMemory) Delete(client string) error {
+	s.mu.Lock()
 
-func (s *InMemory) Close() error {
-	s.cancelFn()
+	if s.data == nil {
+		s.mu.Unlock()
+
+		return nil
+	}
+
+	delete(s.data, client)
+
+	s.mu.Unlock()
 
 	return nil
 }
 
-// collect periodically removes expired tokens from the in-memory store.
-func (s *InMemory) collect(ctx context.Context, ticker *time.Ticker) {
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case <-ticker.C:
-			s.data.Range(func(client, data any) bool {
-				if entry, _ := data.(item); entry.expires.Compare(time.Now()) == -1 {
-					s.data.Delete(client)
-				}
-
-				return true
-			})
-		}
-	}
+// Close is a no-op for in-memory storage, but implements the interface for compatibility.
+func (s *InMemory) Close() error {
+	return nil
 }
diff --git a/internal/tokenstorage/types.go b/internal/tokenstorage/types.go
index d30cc56..e6dfb54 100644
--- a/internal/tokenstorage/types.go
+++ b/internal/tokenstorage/types.go
@@ -1,7 +1,19 @@
 package tokenstorage
 
+import (
+	"time"
+)
+
 type Storage interface {
 	Get(client string) (string, error)
+	Close() error
+	Delete(client string) error
 	Set(client string, token string) error
-	Delete(client string)
 }
+
+type item struct {
+	Expires time.Time
+	Data    []byte
+}
+
+type DataMap map[string]item
diff --git a/packaging/usr/lib/systemd/system/openvpn-auth-oauth2.service b/packaging/usr/lib/systemd/system/openvpn-auth-oauth2.service
index ea1d684..8371d41 100644
--- a/packaging/usr/lib/systemd/system/openvpn-auth-oauth2.service
+++ b/packaging/usr/lib/systemd/system/openvpn-auth-oauth2.service
@@ -1,5 +1,3 @@
-#suppress inspection "DuplicateKeyInSection" for whole file
-
 [Unit]
 Description=OpenVPN authenticator
 Documentation=https://github.com/jkroepke/openvpn-auth-oauth2
@@ -16,8 +14,8 @@ SupplementaryGroups=openvpn-auth-oauth2
 
 ExecStart=/usr/bin/openvpn-auth-oauth2 --config ${CONFIG_FILE}
 NoExecPaths=/
-ExecPaths=/usr/bin/openvpn-auth-oauth2
-ExecReload=kill -HUP $MAINPID
+ExecPaths=/usr/bin/openvpn-auth-oauth2 /usr/bin/env /usr/bin/kill
+ExecReload=/usr/bin/env kill -USR1 $MAINPID
 
 Environment=CONFIG_FILE=/etc/openvpn-auth-oauth2/config.yaml
 EnvironmentFile=-/etc/sysconfig/openvpn-auth-oauth2


Test patch:
diff --git a/internal/openvpn/bench_test.go b/internal/openvpn/bench_test.go
index e5b453b..705ef07 100644
--- a/internal/openvpn/bench_test.go
+++ b/internal/openvpn/bench_test.go
@@ -38,7 +38,7 @@ func BenchmarkOpenVPNHandler(b *testing.B) {
 	conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 	conf.OpenVPN.Bypass = config.OpenVPNBypass{CommonNames: make([]string, 0)}
 
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, b, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(b, managementInterface, openVPNClient)
@@ -104,7 +104,7 @@ func BenchmarkOpenVPNPassthrough(b *testing.B) {
 	conf.OpenVPN.Passthrough.Address = types.URL{URL: &url.URL{Scheme: "tcp", Host: "127.0.0.1:0"}}
 	conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-	tokenStorage := tokenstorage.NewInMemory(b.Context(), testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(b.Context(), b, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(b, managementInterface, openVPNClient)
diff --git a/internal/openvpn/main_test.go b/internal/openvpn/main_test.go
index 6e5e0eb..a7fa7aa 100644
--- a/internal/openvpn/main_test.go
+++ b/internal/openvpn/main_test.go
@@ -42,7 +42,7 @@ func TestClientInvalidServer(t *testing.T) {
 		},
 	}
 
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	err := openVPNClient.Connect(t.Context())
@@ -272,7 +272,7 @@ func TestClientFull(t *testing.T) {
 
 			tc.conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-			tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+			tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 			_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, tc.conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 			managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -371,7 +371,7 @@ func TestClientInvalidPassword(t *testing.T) {
 		},
 	}
 
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -439,7 +439,7 @@ func TestClientInvalidVersion(t *testing.T) {
 			conf.OpenVPN.Bypass = config.OpenVPNBypass{CommonNames: []string{"bypass"}}
 			conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-			tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+			tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 			_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 			managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -491,7 +491,7 @@ func TestHoldRelease(t *testing.T) {
 
 	conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -546,7 +546,7 @@ func TestCommandTimeout(t *testing.T) {
 
 	conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 	_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 	managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -612,7 +612,7 @@ func TestDeadLocks(t *testing.T) {
 
 			conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-			tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+			tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 			_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 			managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
@@ -680,7 +680,7 @@ func TestInvalidCommandResponses(t *testing.T) {
 
 			conf.OpenVPN.Addr = types.URL{URL: &url.URL{Scheme: managementInterface.Addr().Network(), Host: managementInterface.Addr().String()}}
 
-			tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+			tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 			_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 			managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
diff --git a/internal/openvpn/passthrough_test.go b/internal/openvpn/passthrough_test.go
index 4b741a8..cd64e6b 100644
--- a/internal/openvpn/passthrough_test.go
+++ b/internal/openvpn/passthrough_test.go
@@ -200,7 +200,7 @@ func TestPassThroughFull(t *testing.T) {
 				tt.conf.OpenVPN.Passthrough.Address = types.URL{URL: &url.URL{Scheme: tt.scheme, Path: temp}}
 			}
 
-			tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Hour, 5*time.Minute)
+			tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Hour)
 			_, openVPNClient := testutils.SetupOpenVPNOAuth2Clients(ctx, t, tt.conf, logger.Logger, http.DefaultClient, tokenStorage)
 
 			managementInterfaceConn, errOpenVPNClientCh, err := testutils.ConnectToManagementInterface(t, managementInterface, openVPNClient)
diff --git a/internal/tokenstorage/inmemory_test.go b/internal/tokenstorage/inmemory_test.go
index 2ffe216..7db3dfe 100644
--- a/internal/tokenstorage/inmemory_test.go
+++ b/internal/tokenstorage/inmemory_test.go
@@ -1,7 +1,6 @@
 package tokenstorage_test
 
 import (
-	"context"
 	"testing"
 	"time"
 
@@ -11,13 +10,10 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestStorage(t *testing.T) {
+func TestStorage_InMemory(t *testing.T) {
 	t.Parallel()
 
-	ctx, cancel := context.WithCancel(t.Context())
-	t.Cleanup(cancel)
-
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, time.Millisecond*400, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, time.Millisecond*400)
 
 	t.Cleanup(func() {
 		require.NoError(t, tokenStorage.Close())
@@ -44,20 +40,18 @@ func TestStorage(t *testing.T) {
 	_, err = tokenStorage.Get("2")
 	require.ErrorIs(t, err, tokenstorage.ErrNotExists)
 
-	tokenStorage.Delete("1")
+	err = tokenStorage.Delete("1")
+	require.NoError(t, err)
 
 	_, err = tokenStorage.Get("1")
 	require.Error(t, err)
 	require.ErrorIs(t, err, tokenstorage.ErrNotExists)
 }
 
-func TestStorageCleanup(t *testing.T) {
+func TestStorageExpire_InMemory(t *testing.T) {
 	t.Parallel()
 
-	ctx, cancel := context.WithCancel(t.Context())
-	t.Cleanup(cancel)
-
-	tokenStorage := tokenstorage.NewInMemory(ctx, testutils.Secret, 0, time.Millisecond*50)
+	tokenStorage := tokenstorage.NewInMemory(testutils.Secret, 0)
 
 	t.Cleanup(func() {
 		require.NoError(t, tokenStorage.Close())
@@ -66,8 +60,6 @@ func TestStorageCleanup(t *testing.T) {
 	require.NoError(t, tokenStorage.Set("0", "TEST0"))
 	require.NoError(t, tokenStorage.Set("1", "TEST1"))
 
-	time.Sleep(time.Millisecond * 100)
-
 	_, err := tokenStorage.Get("0")
 	require.ErrorIs(t, err, tokenstorage.ErrNotExists)
 
diff --git a/internal/utils/testutils/main.go b/internal/utils/testutils/main.go
index efeafbf..e729125 100644
--- a/internal/utils/testutils/main.go
+++ b/internal/utils/testutils/main.go
@@ -239,7 +239,7 @@ func SetupMockEnvironment(ctx context.Context, tb testing.TB, conf config.Config
 	}
 
 	httpClient := &http.Client{Transport: NewMockRoundTripper(utils.NewUserAgentTransport(rt))}
-	tokenStorage := tokenstorage.NewInMemory(ctx, Secret, conf.OAuth2.Refresh.Expires, 5*time.Minute)
+	tokenStorage := tokenstorage.NewInMemory(Secret, conf.OAuth2.Refresh.Expires)
 
 	oAuth2Client, openvpnClient := SetupOpenVPNOAuth2Clients(ctx, tb, conf, logger.Logger, httpClient, tokenStorage)
 
diff --git a/internal/utils/testutils/storage.go b/internal/utils/testutils/storage.go
index 23a7ffb..8b56832 100644
--- a/internal/utils/testutils/storage.go
+++ b/internal/utils/testutils/storage.go
@@ -10,9 +10,14 @@ func (FakeStorage) Get(_ string) (string, error) {
 	return "", nil
 }
 
-func (FakeStorage) Set(_, _ string) error {
+func (FakeStorage) Close() error {
+	return nil
+}
+
+func (FakeStorage) Delete(_ string) error {
 	return nil
 }
 
-func (FakeStorage) Delete(_ string) {
+func (FakeStorage) Set(_, _ string) error {
+	return nil
 }