Fix startup failure when remote identifier is not present

Signed-off-by: Stefano Scafiti <stefano.scafiti96@gmail.com>
codenotary · Jul 24, 2024 · 3eec9d4 · 3eec9d4
1 parent 10ad005
commit 3eec9d4
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 48 deletions.
diff --git a/pkg/server/remote_storage.go b/pkg/server/remote_storage.go
@@ -19,7 +19,7 @@ package server
 import (
 	"bytes"
 	"context"
-	"io/ioutil"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
@@ -75,67 +75,82 @@ func (s *ImmuServer) initializeRemoteStorage(storage remotestorage.Storage) erro
 		return nil
 	}
 
-	ctx := context.Background()
+	if s.Options.RemoteStorageOptions.S3ExternalIdentifier {
+		if err := s.loadRemoteIdentifier(context.Background(), storage); err != nil {
+			return err
+		}
+	}
+	return s.createRemoteSubFolders(storage)
+}
 
-	hasRemoteIdentifier, err := storage.Exists(ctx, IDENTIFIER_FNAME)
+func (s *ImmuServer) createRemoteSubFolders(storage remotestorage.Storage) error {
+	_, subFolders, err := storage.ListEntries(context.Background(), "")
 	if err != nil {
 		return err
 	}
 
-	if !hasRemoteIdentifier && s.Options.RemoteStorageOptions.S3ExternalIdentifier {
-		return ErrNoRemoteIdentifier
-	}
-
-	localIdentifierFile := filepath.Join(s.Options.Dir, IDENTIFIER_FNAME)
-
-	if hasRemoteIdentifier {
-		remoteIDStream, err := storage.Get(ctx, IDENTIFIER_FNAME, 0, -1)
-		if err != nil {
-			return err
-		}
-		remoteID, err := ioutil.ReadAll(remoteIDStream)
-		remoteIDStream.Close()
+	for _, subFolder := range subFolders {
+		err := os.MkdirAll(
+			filepath.Join(s.Options.Dir, subFolder),
+			store.DefaultFileMode,
+		)
 		if err != nil {
 			return err
 		}
+	}
+	return nil
+}
 
-		if !fileExists(localIdentifierFile) {
-			err := ioutil.WriteFile(localIdentifierFile, remoteID, os.ModePerm)
-			if err != nil {
-				return err
-			}
-			s.UUID, err = xid.FromBytes(remoteID)
-			if err != nil {
-				return err
-			}
-		} else {
-			localID, err := ioutil.ReadFile(localIdentifierFile)
-			if err != nil {
-				return err
-			}
+func (s *ImmuServer) loadRemoteIdentifier(ctx context.Context, storage remotestorage.Storage) error {
+	hasRemoteIdentifier, err := storage.Exists(ctx, IDENTIFIER_FNAME)
+	if err != nil {
+		return err
+	}
 
-			if !bytes.Equal(remoteID, localID) {
-				return ErrRemoteStorageDoesNotMatch
-			}
-		}
+	if !hasRemoteIdentifier {
+		return s.initRemoteIdentifier(ctx, storage)
 	}
 
-	// Ensure all sub-folders are created, init code relies on this
-	_, subFolders, err := storage.ListEntries(context.Background(), "")
+	remoteIDStream, err := storage.Get(ctx, IDENTIFIER_FNAME, 0, -1)
 	if err != nil {
 		return err
 	}
-	for _, subFolder := range subFolders {
-		err := os.MkdirAll(
-			filepath.Join(s.Options.Dir, subFolder),
-			store.DefaultFileMode,
-		)
+
+	remoteID, err := io.ReadAll(remoteIDStream)
+	remoteIDStream.Close()
+	if err != nil {
+		return err
+	}
+
+	localIdentifierFile := filepath.Join(s.Options.Dir, IDENTIFIER_FNAME)
+	if fileExists(localIdentifierFile) {
+		localID, err := os.ReadFile(localIdentifierFile)
 		if err != nil {
 			return err
 		}
+
+		if !bytes.Equal(remoteID, localID) {
+			return ErrRemoteStorageDoesNotMatch
+		}
+		return nil
 	}
 
-	return nil
+	if err := os.WriteFile(localIdentifierFile, remoteID, os.ModePerm); err != nil {
+		return err
+	}
+
+	s.UUID, err = xid.FromBytes(remoteID)
+	return err
+}
+
+func (s *ImmuServer) initRemoteIdentifier(ctx context.Context, storage remotestorage.Storage) error {
+	localIdentifierFile := filepath.Join(s.Options.Dir, IDENTIFIER_FNAME)
+
+	s.UUID = xid.New()
+	if err := os.WriteFile(localIdentifierFile, s.UUID.Bytes(), os.ModePerm); err != nil {
+		return err
+	}
+	return storage.Put(ctx, IDENTIFIER_FNAME, localIdentifierFile)
 }
 
 func (s *ImmuServer) updateRemoteUUID(remoteStorage remotestorage.Storage) error {

diff --git a/pkg/server/remote_storage_test.go b/pkg/server/remote_storage_test.go
@@ -175,7 +175,8 @@ func TestInitializeRemoteStorageEmptyRemoteStorage(t *testing.T) {
 func TestInitializeRemoteStorageEmptyRemoteStorageErrorOnExists(t *testing.T) {
 	dir := t.TempDir()
 
-	opts := DefaultOptions().WithDir(dir)
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithDir(dir).WithRemoteStorageOptions(remoteStorageOpts)
 
 	s := DefaultServer()
 
@@ -190,7 +191,7 @@ func TestInitializeRemoteStorageEmptyRemoteStorageErrorOnExists(t *testing.T) {
 	}
 
 	err := s.initializeRemoteStorage(mem)
-	require.True(t, errors.Is(err, injectedErr))
+	require.ErrorIs(t, err, injectedErr)
 }
 
 func TestInitializeRemoteStorageEmptyRemoteStorageErrorOnListEntries(t *testing.T) {
@@ -217,7 +218,8 @@ func TestInitializeRemoteStorageEmptyRemoteStorageErrorOnListEntries(t *testing.
 func TestInitializeRemoteStorageDownloadIdentifier(t *testing.T) {
 	dir := t.TempDir()
 
-	opts := DefaultOptions().WithDir(dir)
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithDir(dir).WithRemoteStorageOptions(remoteStorageOpts)
 
 	s := DefaultServer()
 
@@ -301,7 +303,8 @@ func TestInitializeRemoteStorageWithoutLocalIdentifier(t *testing.T) {
 func TestInitializeRemoteStorageDownloadIdentifierErrorOnGet(t *testing.T) {
 	dir := t.TempDir()
 
-	opts := DefaultOptions().WithDir(dir)
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithDir(dir).WithRemoteStorageOptions(remoteStorageOpts)
 
 	s := DefaultServer()
 
@@ -324,7 +327,12 @@ func TestInitializeRemoteStorageDownloadIdentifierErrorOnGet(t *testing.T) {
 func TestInitializeRemoteStorageDownloadIdentifierErrorOnStore(t *testing.T) {
 	require.NoError(t, os.MkdirAll(filepath.Join(t.TempDir(), "data_uuiderr", "immudb.identifier"), 0777))
 
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithRemoteStorageOptions(remoteStorageOpts)
+
 	s := DefaultServer()
+	s.WithOptions(opts)
+
 	m := memory.Open()
 	storeData(t, m, "immudb.identifier", []byte{1, 2, 3, 4, 5})
 	err := s.initializeRemoteStorage(m)
@@ -342,7 +350,8 @@ func (e errReader) Read([]byte) (int, error) {
 func TestInitializeRemoteStorageDownloadIdentifierErrorOnRead(t *testing.T) {
 	dir := t.TempDir()
 
-	opts := DefaultOptions().WithDir(dir)
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithDir(dir).WithRemoteStorageOptions(remoteStorageOpts)
 
 	s := DefaultServer()
 
@@ -365,7 +374,8 @@ func TestInitializeRemoteStorageDownloadIdentifierErrorOnRead(t *testing.T) {
 func TestInitializeRemoteStorageIdentifierMismatch(t *testing.T) {
 	dir := t.TempDir()
 
-	opts := DefaultOptions().WithDir(dir)
+	remoteStorageOpts := DefaultRemoteStorageOptions().WithS3ExternalIdentifier(true)
+	opts := DefaultOptions().WithDir(dir).WithRemoteStorageOptions(remoteStorageOpts)
 
 	s := DefaultServer()