internal/databroker: store server version (#1121)

Storing server version when creating new server. After then, we can
retrieve the version from backend when server restart.

With storage backend which supports persistent, the server version
won't change after restarting.

internal/databroker/server.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/golang/protobuf/ptypes"
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
@@ -22,6 +23,11 @@ import (
 	"github.com/pomerium/pomerium/pkg/storage/inmemory"
 )
 
+const (
+	recordTypeServerVersion = "server_version"
+	serverVersionKey        = "version"
+)
+
 // Server implements the databroker service using an in memory database.
 type Server struct {
 	version string
@@ -44,6 +50,8 @@ func New(options ...ServerOption) *Server {
 		byType:   make(map[string]storage.Backend),
 		onchange: NewSignal(),
 	}
+	srv.initVersion()
+
 	go func() {
 		ticker := time.NewTicker(cfg.deletePermanentlyAfter / 2)
 		defer ticker.Stop()
@@ -64,6 +72,28 @@ func New(options ...ServerOption) *Server {
 	return srv
 }
 
+func (srv *Server) initVersion() {
+	dbServerVersion := srv.getDB(recordTypeServerVersion)
+	if dbServerVersion == nil {
+		return
+	}
+
+	// Get version from storage first.
+	if r := dbServerVersion.Get(context.Background(), serverVersionKey); r != nil {
+		var sv databroker.ServerVersion
+		if err := ptypes.UnmarshalAny(r.GetData(), &sv); err != nil {
+			srv.log.Debug().Str("server_version", sv.Version).Msg("got db version from DB")
+			srv.version = sv.Version
+		}
+		return
+	}
+
+	data, _ := ptypes.MarshalAny(&databroker.ServerVersion{Version: srv.version})
+	if err := dbServerVersion.Put(context.Background(), serverVersionKey, data); err != nil {
+		srv.log.Warn().Err(err).Msg("failed to save server version.")
+	}
+}
+
 // Delete deletes a record from the in-memory list.
 func (srv *Server) Delete(ctx context.Context, req *databroker.DeleteRequest) (*empty.Empty, error) {
 	_, span := trace.StartSpan(ctx, "databroker.grpc.Delete")

internal/databroker/server_test.go

@@ -0,0 +1,73 @@
+package databroker
+
+import (
+	"context"
+	"testing"
+
+	"github.com/golang/protobuf/ptypes"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/pomerium/pomerium/internal/log"
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/storage"
+)
+
+func newServer(cfg *serverConfig) *Server {
+	return &Server{
+		version: uuid.New().String(),
+		cfg:     cfg,
+		log:     log.With().Str("service", "databroker").Logger(),
+
+		byType:   make(map[string]storage.Backend),
+		onchange: NewSignal(),
+	}
+}
+
+func TestServer_initVersion(t *testing.T) {
+	cfg := newServerConfig()
+	t.Run("nil db", func(t *testing.T) {
+		srv := newServer(cfg)
+		srvVersion := uuid.New().String()
+		srv.version = srvVersion
+		srv.byType[recordTypeServerVersion] = nil
+		srv.initVersion()
+		assert.Equal(t, srvVersion, srv.version)
+	})
+	t.Run("new server with random version", func(t *testing.T) {
+		srv := newServer(cfg)
+		ctx := context.Background()
+		db := srv.getDB(recordTypeServerVersion)
+		r := db.Get(ctx, serverVersionKey)
+		assert.Nil(t, r)
+		srvVersion := uuid.New().String()
+		srv.version = srvVersion
+		srv.initVersion()
+		assert.Equal(t, srvVersion, srv.version)
+		r = db.Get(ctx, serverVersionKey)
+		assert.NotNil(t, r)
+		var sv databroker.ServerVersion
+		assert.NoError(t, ptypes.UnmarshalAny(r.GetData(), &sv))
+		assert.Equal(t, srv.version, sv.Version)
+	})
+	t.Run("init version twice should get the same version", func(t *testing.T) {
+		srv := newServer(cfg)
+		ctx := context.Background()
+		db := srv.getDB(recordTypeServerVersion)
+		r := db.Get(ctx, serverVersionKey)
+		assert.Nil(t, r)
+		srvVersion := uuid.New().String()
+		srv.version = srvVersion
+		srv.initVersion()
+		assert.Equal(t, srvVersion, srv.version)
+		r = db.Get(ctx, serverVersionKey)
+		assert.NotNil(t, r)
+		var sv databroker.ServerVersion
+		assert.NoError(t, ptypes.UnmarshalAny(r.GetData(), &sv))
+		assert.Equal(t, srv.version, sv.Version)
+
+		// re-init version should get the same value as above
+		srv.initVersion()
+		assert.Equal(t, srvVersion, srv.version)
+	})
+}