fully migrate DNS to proto

BeryJu · May 31, 2023 · 7f31332 · 7f31332
1 parent cbaebfa
commit 7f31332
Show file tree

Hide file tree

Showing 9 changed files with 192 additions and 32 deletions.
diff --git a/pkg/roles/dns/record.go b/pkg/roles/dns/record.go
@@ -2,12 +2,12 @@ package dns
 
 import (
 	"context"
-	"encoding/json"
 	"net"
 	"strings"
 
 	"beryju.io/gravity/pkg/roles"
 	"beryju.io/gravity/pkg/roles/dns/types"
+	"beryju.io/gravity/pkg/storage"
 	"github.com/miekg/dns"
 	"go.etcd.io/etcd/api/v3/mvccpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
@@ -43,13 +43,7 @@ func (z *ZoneContext) recordFromKV(kv *mvccpb.KeyValue) (*RecordContext, error)
 	}
 	rec.recordKey = strings.TrimSuffix(fullRecordKey, "/"+rec.uid)
 
-	// Try loading protobuf first
-	err := proto.Unmarshal(kv.Value, rec.Record)
-	if err == nil {
-		return rec, nil
-	}
-	// Otherwise try json
-	err = json.Unmarshal(kv.Value, &rec)
+	_, err := storage.Parse(kv.Value, rec.Record)
 	if err != nil {
 		return rec, err
 	}

diff --git a/pkg/roles/dns/role.go b/pkg/roles/dns/role.go
@@ -24,7 +24,7 @@ type Role struct {
 	ctx   context.Context
 	zones map[string]*ZoneContext
 
-	cfg     *RoleConfig
+	cfg     *types.RoleConfig
 	log     *zap.Logger
 	servers []*dns.Server
 	zonesM  sync.RWMutex

diff --git a/pkg/roles/dns/role_config.go b/pkg/roles/dns/role_config.go
@@ -2,41 +2,38 @@ package dns
 
 import (
 	"context"
-	"encoding/json"
 
 	instanceTypes "beryju.io/gravity/pkg/instance/types"
+	"beryju.io/gravity/pkg/storage"
 	"go.uber.org/zap"
+	"google.golang.org/protobuf/proto"
 
 	"beryju.io/gravity/pkg/roles/dns/types"
 	"github.com/swaggest/usecase"
 	"github.com/swaggest/usecase/status"
 )
 
-type RoleConfig struct {
-	Port int32 `json:"port"`
-}
-
-func (r *Role) decodeRoleConfig(raw []byte) *RoleConfig {
-	def := RoleConfig{
+func (r *Role) decodeRoleConfig(raw []byte) *types.RoleConfig {
+	def := types.RoleConfig{
 		Port: 53,
 	}
 	if len(raw) < 1 {
 		return &def
 	}
-	err := json.Unmarshal(raw, &def)
+	conf, err := storage.Parse(raw, &def)
 	if err != nil {
 		r.log.Warn("failed to decode role config", zap.Error(err))
 	}
-	return &def
+	return conf
 }
 
 type APIRoleConfigOutput struct {
-	Config RoleConfig `json:"config" required:"true"`
+	Config *types.RoleConfig `json:"config" required:"true"`
 }
 
 func (r *Role) APIRoleConfigGet() usecase.Interactor {
 	u := usecase.NewInteractor(func(ctx context.Context, input struct{}, output *APIRoleConfigOutput) error {
-		output.Config = *r.cfg
+		output.Config = r.cfg
 		return nil
 	})
 	u.SetName("dns.get_role_config")
@@ -46,12 +43,12 @@ func (r *Role) APIRoleConfigGet() usecase.Interactor {
 }
 
 type APIRoleConfigInput struct {
-	Config RoleConfig `json:"config" required:"true"`
+	Config *types.RoleConfig `json:"config" required:"true"`
 }
 
 func (r *Role) APIRoleConfigPut() usecase.Interactor {
 	u := usecase.NewInteractor(func(ctx context.Context, input APIRoleConfigInput, output *struct{}) error {
-		jc, err := json.Marshal(input.Config)
+		jc, err := proto.Marshal(input.Config)
 		if err != nil {
 			return status.Wrap(err, status.InvalidArgument)
 		}

diff --git a/pkg/roles/dns/role_config_test.go b/pkg/roles/dns/role_config_test.go
@@ -5,6 +5,7 @@ import (
 
 	"beryju.io/gravity/pkg/instance"
 	"beryju.io/gravity/pkg/roles/dns"
+	"beryju.io/gravity/pkg/roles/dns/types"
 	"beryju.io/gravity/pkg/tests"
 	"github.com/stretchr/testify/assert"
 )
@@ -31,7 +32,7 @@ func TestAPIRoleConfigPut(t *testing.T) {
 	defer role.Stop()
 
 	assert.NoError(t, role.APIRoleConfigPut().Interact(ctx, dns.APIRoleConfigInput{
-		Config: dns.RoleConfig{
+		Config: &types.RoleConfig{
 			Port: 1054,
 		},
 	}, &struct{}{}))

diff --git a/pkg/roles/dns/role_test.go b/pkg/roles/dns/role_test.go
@@ -5,12 +5,13 @@ import (
 
 	"beryju.io/gravity/pkg/instance"
 	"beryju.io/gravity/pkg/roles/dns"
+	"beryju.io/gravity/pkg/roles/dns/types"
 	"beryju.io/gravity/pkg/tests"
 	"github.com/stretchr/testify/assert"
 )
 
 func RoleConfig() []byte {
-	return []byte(tests.MustJSON(dns.RoleConfig{
+	return []byte(tests.MustPB(&types.RoleConfig{
 		Port: 1054,
 	}))
 }

diff --git a/pkg/roles/dns/types/role_dns_config.pb.go b/pkg/roles/dns/types/role_dns_config.pb.go
diff --git a/pkg/roles/dns/zone.go b/pkg/roles/dns/zone.go
@@ -2,7 +2,6 @@ package dns
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"strings"
 	"sync"
@@ -12,6 +11,7 @@ import (
 	"beryju.io/gravity/pkg/roles/dns/types"
 	"beryju.io/gravity/pkg/roles/dns/utils"
 	tsdbTypes "beryju.io/gravity/pkg/roles/tsdb/types"
+	"beryju.io/gravity/pkg/storage"
 	"github.com/getsentry/sentry-go"
 	"github.com/miekg/dns"
 	"go.etcd.io/etcd/api/v3/mvccpb"
@@ -139,13 +139,7 @@ func (r *Role) zoneFromKV(raw *mvccpb.KeyValue) (*ZoneContext, error) {
 	z.log = r.log.With(zap.String("zone", name))
 	z.etcdKey = string(raw.Key)
 
-	// Try loading protobuf first
-	err := proto.Unmarshal(raw.Value, z.Zone)
-	if err == nil {
-		return z, nil
-	}
-	// Otherwise try json
-	err = json.Unmarshal(raw.Value, &z)
+	_, err := storage.Parse(raw.Value, z.Zone)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/storage/utils.go b/pkg/storage/utils.go
@@ -0,0 +1,23 @@
+package storage
+
+import (
+	"encoding/json"
+
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+// Parse Attempt to parse data as protobuf first, otherwise try json
+func Parse[T protoreflect.ProtoMessage](raw []byte, v T) (T, error) {
+	// Try loading protobuf first
+	err := proto.Unmarshal(raw, v)
+	if err == nil {
+		return v, nil
+	}
+	// Otherwise try json
+	err = json.Unmarshal(raw, &v)
+	if err == nil {
+		return v, nil
+	}
+	return v, err
+}
diff --git a/protobuf/role_dns_config.proto b/protobuf/role_dns_config.proto
@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+option go_package = "pkg/roles/dns/types";
+
+message RoleConfig {
+    int32 port = 1;
+}