From abff3fa57c6f2a7b429a13ffd9ba0b5fffe9d195 Mon Sep 17 00:00:00 2001 From: Arthur Bellal Date: Fri, 7 Jan 2022 21:28:56 +0100 Subject: [PATCH] (rcm-40) improve key format (#10278) * (rcm-40) improve key format * remove proto key * add msgpkey * migrate to msgp * fix linting issues * python lint --- pkg/config/remote/service/service.go | 8 +- pkg/config/remote/service/service_test.go | 13 +- pkg/config/remote/service/util.go | 33 ++--- pkg/config/remote/service/util_test.go | 46 +++++++ pkg/proto/msgpgo/key.go | 8 ++ pkg/proto/msgpgo/key_gen.go | 160 ++++++++++++++++++++++ pkg/proto/msgpgo/key_gen_test.go | 123 +++++++++++++++++ pkg/proto/pbgo/mocks/api_mockgen.pb.go | 2 +- tasks/go.py | 3 + 9 files changed, 368 insertions(+), 28 deletions(-) create mode 100644 pkg/config/remote/service/util_test.go create mode 100644 pkg/proto/msgpgo/key.go create mode 100644 pkg/proto/msgpgo/key_gen.go create mode 100644 pkg/proto/msgpgo/key_gen_test.go diff --git a/pkg/config/remote/service/service.go b/pkg/config/remote/service/service.go index c94e1d29193e4..5f0b3e7f191a5 100644 --- a/pkg/config/remote/service/service.go +++ b/pkg/config/remote/service/service.go @@ -38,7 +38,6 @@ type Service struct { firstUpdate bool refreshInterval time.Duration - remoteConfigKey remoteConfigKey ctx context.Context clock clock.Clock @@ -86,15 +85,15 @@ func NewService() (*Service, error) { return nil, err } backendURL := config.Datadog.GetString("remote_configuration.endpoint") - http := api.NewHTTPClient(backendURL, apiKey, remoteConfigKey.appKey) + http := api.NewHTTPClient(backendURL, apiKey, remoteConfigKey.AppKey) dbPath := path.Join(config.Datadog.GetString("run_path"), "remote-config.db") db, err := openCacheDB(dbPath) if err != nil { return nil, err } - cacheKey := fmt.Sprintf("%s/%d/", remoteConfigKey.datacenter, remoteConfigKey.orgID) - uptaneClient, err := uptane.NewClient(db, cacheKey, remoteConfigKey.orgID) + cacheKey := fmt.Sprintf("%s/%d/", remoteConfigKey.Datacenter, remoteConfigKey.OrgID) + uptaneClient, err := uptane.NewClient(db, cacheKey, remoteConfigKey.OrgID) if err != nil { return nil, err } @@ -109,7 +108,6 @@ func NewService() (*Service, error) { ctx: context.Background(), firstUpdate: true, refreshInterval: refreshInterval, - remoteConfigKey: remoteConfigKey, products: make(map[rdata.Product]struct{}), newProducts: make(map[rdata.Product]struct{}), hostname: hostname, diff --git a/pkg/config/remote/service/service_test.go b/pkg/config/remote/service/service_test.go index 70d7082cd351f..604b6c44e0ae8 100644 --- a/pkg/config/remote/service/service_test.go +++ b/pkg/config/remote/service/service_test.go @@ -2,12 +2,14 @@ package service import ( "context" + "encoding/base32" "os" "testing" "github.com/DataDog/datadog-agent/pkg/config" rdata "github.com/DataDog/datadog-agent/pkg/config/remote/data" "github.com/DataDog/datadog-agent/pkg/config/remote/uptane" + "github.com/DataDog/datadog-agent/pkg/proto/msgpgo" "github.com/DataDog/datadog-agent/pkg/proto/pbgo" "github.com/DataDog/datadog-agent/pkg/version" "github.com/benbjohnson/clock" @@ -59,15 +61,20 @@ func (m *mockUptane) TargetsMeta() ([]byte, error) { return args.Get(0).([]byte), args.Error(1) } -const ( - testRCKey = "dd.com/2/fake_key" +var ( + testRCKey = msgpgo.RemoteConfigKey{ + AppKey: "fake_key", + OrgID: 2, + Datacenter: "dd.com", + } ) func newTestService(t *testing.T, api *mockAPI, uptane *mockUptane, clock clock.Clock) *Service { dir, err := os.MkdirTemp("", "testdbdir") assert.NoError(t, err) config.Datadog.Set("run_path", dir) - config.Datadog.Set("remote_configuration.key", testRCKey) + serializedKey, _ := testRCKey.MarshalMsg(nil) + config.Datadog.Set("remote_configuration.key", base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(serializedKey)) service, err := NewService() assert.NoError(t, err) service.api = api diff --git a/pkg/config/remote/service/util.go b/pkg/config/remote/service/util.go index 4b1544d560def..f5f6c564338dc 100644 --- a/pkg/config/remote/service/util.go +++ b/pkg/config/remote/service/util.go @@ -6,13 +6,13 @@ package service import ( + "encoding/base32" "fmt" "os" - "strconv" - "strings" "github.com/DataDog/datadog-agent/pkg/config/remote/data" "github.com/DataDog/datadog-agent/pkg/config/remote/uptane" + "github.com/DataDog/datadog-agent/pkg/proto/msgpgo" "github.com/DataDog/datadog-agent/pkg/proto/pbgo" "github.com/DataDog/datadog-agent/pkg/version" "go.etcd.io/bbolt" @@ -31,26 +31,21 @@ func openCacheDB(path string) (*bbolt.DB, error) { return db, nil } -type remoteConfigKey struct { - orgID int64 - appKey string - datacenter string -} - -func parseRemoteConfigKey(rawKey string) (remoteConfigKey, error) { - split := strings.SplitN(rawKey, "/", 3) - if len(split) < 3 { - return remoteConfigKey{}, fmt.Errorf("invalid remote configuration key format, should be datacenter/org_id/app_key") +func parseRemoteConfigKey(serializedKey string) (*msgpgo.RemoteConfigKey, error) { + encoding := base32.StdEncoding.WithPadding(base32.NoPadding) + rawKey, err := encoding.DecodeString(serializedKey) + if err != nil { + return nil, err } - orgID, err := strconv.ParseInt(split[1], 10, 64) + var key msgpgo.RemoteConfigKey + _, err = key.UnmarshalMsg(rawKey) if err != nil { - return remoteConfigKey{}, err + return nil, err + } + if key.AppKey == "" || key.Datacenter == "" || key.OrgID == 0 { + return nil, fmt.Errorf("invalid remote config key") } - return remoteConfigKey{ - orgID: orgID, - appKey: split[2], - datacenter: split[0], - }, nil + return &key, nil } func buildLatestConfigsRequest(hostname string, state uptane.State, activeClients []*pbgo.Client, products map[data.Product]struct{}, newProducts map[data.Product]struct{}) *pbgo.LatestConfigsRequest { diff --git a/pkg/config/remote/service/util_test.go b/pkg/config/remote/service/util_test.go new file mode 100644 index 0000000000000..c3766a263574e --- /dev/null +++ b/pkg/config/remote/service/util_test.go @@ -0,0 +1,46 @@ +package service + +import ( + "encoding/base32" + "testing" + + "github.com/DataDog/datadog-agent/pkg/proto/msgpgo" + "github.com/stretchr/testify/assert" +) + +func TestRemoteConfigKey(t *testing.T) { + tests := []struct { + input string + err bool + output *msgpgo.RemoteConfigKey + }{ + {input: generateKey(t, 2, "datadoghq.com", "58d58c60b8ac337293ce2ca6b28b19eb"), output: &msgpgo.RemoteConfigKey{AppKey: "58d58c60b8ac337293ce2ca6b28b19eb", OrgID: 2, Datacenter: "datadoghq.com"}}, + {input: generateKey(t, 2, "datadoghq.com", ""), err: true}, + {input: generateKey(t, 2, "", "app_Key"), err: true}, + {input: generateKey(t, 0, "datadoghq.com", "app_Key"), err: true}, + } + for _, test := range tests { + t.Run(test.input, func(tt *testing.T) { + output, err := parseRemoteConfigKey(test.input) + if test.err { + assert.Error(tt, err) + } else { + assert.Equal(tt, test.output, output) + assert.NoError(tt, err) + } + }) + } +} + +func generateKey(t *testing.T, orgID int64, datacenter string, appKey string) string { + key := msgpgo.RemoteConfigKey{ + AppKey: appKey, + OrgID: orgID, + Datacenter: datacenter, + } + rawKey, err := key.MarshalMsg(nil) + if err != nil { + t.Fatal(err) + } + return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(rawKey) +} diff --git a/pkg/proto/msgpgo/key.go b/pkg/proto/msgpgo/key.go new file mode 100644 index 0000000000000..ea894da31e270 --- /dev/null +++ b/pkg/proto/msgpgo/key.go @@ -0,0 +1,8 @@ +package msgpgo + +// RemoteConfigKey is a RemoteConfigKey +type RemoteConfigKey struct { + AppKey string `msgpack:"key"` + OrgID int64 `msgpack:"org"` + Datacenter string `msgpack:"dc"` +} diff --git a/pkg/proto/msgpgo/key_gen.go b/pkg/proto/msgpgo/key_gen.go new file mode 100644 index 0000000000000..b85a3cedd3689 --- /dev/null +++ b/pkg/proto/msgpgo/key_gen.go @@ -0,0 +1,160 @@ +package msgpgo + +// Code generated by github.com/tinylib/msgp DO NOT EDIT. + +import ( + "github.com/tinylib/msgp/msgp" +) + +// DecodeMsg implements msgp.Decodable +func (z *RemoteConfigKey) DecodeMsg(dc *msgp.Reader) (err error) { + var field []byte + _ = field + var zb0001 uint32 + zb0001, err = dc.ReadMapHeader() + if err != nil { + err = msgp.WrapError(err) + return + } + for zb0001 > 0 { + zb0001-- + field, err = dc.ReadMapKeyPtr() + if err != nil { + err = msgp.WrapError(err) + return + } + switch msgp.UnsafeString(field) { + case "key": + z.AppKey, err = dc.ReadString() + if err != nil { + err = msgp.WrapError(err, "AppKey") + return + } + case "org": + z.OrgID, err = dc.ReadInt64() + if err != nil { + err = msgp.WrapError(err, "OrgID") + return + } + case "dc": + z.Datacenter, err = dc.ReadString() + if err != nil { + err = msgp.WrapError(err, "Datacenter") + return + } + default: + err = dc.Skip() + if err != nil { + err = msgp.WrapError(err) + return + } + } + } + return +} + +// EncodeMsg implements msgp.Encodable +func (z RemoteConfigKey) EncodeMsg(en *msgp.Writer) (err error) { + // map header, size 3 + // write "key" + err = en.Append(0x83, 0xa3, 0x6b, 0x65, 0x79) + if err != nil { + return + } + err = en.WriteString(z.AppKey) + if err != nil { + err = msgp.WrapError(err, "AppKey") + return + } + // write "org" + err = en.Append(0xa3, 0x6f, 0x72, 0x67) + if err != nil { + return + } + err = en.WriteInt64(z.OrgID) + if err != nil { + err = msgp.WrapError(err, "OrgID") + return + } + // write "dc" + err = en.Append(0xa2, 0x64, 0x63) + if err != nil { + return + } + err = en.WriteString(z.Datacenter) + if err != nil { + err = msgp.WrapError(err, "Datacenter") + return + } + return +} + +// MarshalMsg implements msgp.Marshaler +func (z RemoteConfigKey) MarshalMsg(b []byte) (o []byte, err error) { + o = msgp.Require(b, z.Msgsize()) + // map header, size 3 + // string "key" + o = append(o, 0x83, 0xa3, 0x6b, 0x65, 0x79) + o = msgp.AppendString(o, z.AppKey) + // string "org" + o = append(o, 0xa3, 0x6f, 0x72, 0x67) + o = msgp.AppendInt64(o, z.OrgID) + // string "dc" + o = append(o, 0xa2, 0x64, 0x63) + o = msgp.AppendString(o, z.Datacenter) + return +} + +// UnmarshalMsg implements msgp.Unmarshaler +func (z *RemoteConfigKey) UnmarshalMsg(bts []byte) (o []byte, err error) { + var field []byte + _ = field + var zb0001 uint32 + zb0001, bts, err = msgp.ReadMapHeaderBytes(bts) + if err != nil { + err = msgp.WrapError(err) + return + } + for zb0001 > 0 { + zb0001-- + field, bts, err = msgp.ReadMapKeyZC(bts) + if err != nil { + err = msgp.WrapError(err) + return + } + switch msgp.UnsafeString(field) { + case "key": + z.AppKey, bts, err = msgp.ReadStringBytes(bts) + if err != nil { + err = msgp.WrapError(err, "AppKey") + return + } + case "org": + z.OrgID, bts, err = msgp.ReadInt64Bytes(bts) + if err != nil { + err = msgp.WrapError(err, "OrgID") + return + } + case "dc": + z.Datacenter, bts, err = msgp.ReadStringBytes(bts) + if err != nil { + err = msgp.WrapError(err, "Datacenter") + return + } + default: + bts, err = msgp.Skip(bts) + if err != nil { + err = msgp.WrapError(err) + return + } + } + } + o = bts + return +} + +// Msgsize returns an upper bound estimate of the number of bytes occupied by the serialized message +func (z RemoteConfigKey) Msgsize() (s int) { + s = 1 + 4 + msgp.StringPrefixSize + len(z.AppKey) + 4 + msgp.Int64Size + 3 + msgp.StringPrefixSize + len(z.Datacenter) + return +} diff --git a/pkg/proto/msgpgo/key_gen_test.go b/pkg/proto/msgpgo/key_gen_test.go new file mode 100644 index 0000000000000..30c16604824ad --- /dev/null +++ b/pkg/proto/msgpgo/key_gen_test.go @@ -0,0 +1,123 @@ +package msgpgo + +// Code generated by github.com/tinylib/msgp DO NOT EDIT. + +import ( + "bytes" + "testing" + + "github.com/tinylib/msgp/msgp" +) + +func TestMarshalUnmarshalRemoteConfigKey(t *testing.T) { + v := RemoteConfigKey{} + bts, err := v.MarshalMsg(nil) + if err != nil { + t.Fatal(err) + } + left, err := v.UnmarshalMsg(bts) + if err != nil { + t.Fatal(err) + } + if len(left) > 0 { + t.Errorf("%d bytes left over after UnmarshalMsg(): %q", len(left), left) + } + + left, err = msgp.Skip(bts) + if err != nil { + t.Fatal(err) + } + if len(left) > 0 { + t.Errorf("%d bytes left over after Skip(): %q", len(left), left) + } +} + +func BenchmarkMarshalMsgRemoteConfigKey(b *testing.B) { + v := RemoteConfigKey{} + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + v.MarshalMsg(nil) + } +} + +func BenchmarkAppendMsgRemoteConfigKey(b *testing.B) { + v := RemoteConfigKey{} + bts := make([]byte, 0, v.Msgsize()) + bts, _ = v.MarshalMsg(bts[0:0]) + b.SetBytes(int64(len(bts))) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + bts, _ = v.MarshalMsg(bts[0:0]) + } +} + +func BenchmarkUnmarshalRemoteConfigKey(b *testing.B) { + v := RemoteConfigKey{} + bts, _ := v.MarshalMsg(nil) + b.ReportAllocs() + b.SetBytes(int64(len(bts))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := v.UnmarshalMsg(bts) + if err != nil { + b.Fatal(err) + } + } +} + +func TestEncodeDecodeRemoteConfigKey(t *testing.T) { + v := RemoteConfigKey{} + var buf bytes.Buffer + msgp.Encode(&buf, &v) + + m := v.Msgsize() + if buf.Len() > m { + t.Log("WARNING: TestEncodeDecodeRemoteConfigKey Msgsize() is inaccurate") + } + + vn := RemoteConfigKey{} + err := msgp.Decode(&buf, &vn) + if err != nil { + t.Error(err) + } + + buf.Reset() + msgp.Encode(&buf, &v) + err = msgp.NewReader(&buf).Skip() + if err != nil { + t.Error(err) + } +} + +func BenchmarkEncodeRemoteConfigKey(b *testing.B) { + v := RemoteConfigKey{} + var buf bytes.Buffer + msgp.Encode(&buf, &v) + b.SetBytes(int64(buf.Len())) + en := msgp.NewWriter(msgp.Nowhere) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + v.EncodeMsg(en) + } + en.Flush() +} + +func BenchmarkDecodeRemoteConfigKey(b *testing.B) { + v := RemoteConfigKey{} + var buf bytes.Buffer + msgp.Encode(&buf, &v) + b.SetBytes(int64(buf.Len())) + rd := msgp.NewEndlessReader(buf.Bytes(), b) + dc := msgp.NewReader(rd) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + err := v.DecodeMsg(dc) + if err != nil { + b.Fatal(err) + } + } +} diff --git a/pkg/proto/pbgo/mocks/api_mockgen.pb.go b/pkg/proto/pbgo/mocks/api_mockgen.pb.go index 2f2adf9fc4287..5e799d697b0ef 100644 --- a/pkg/proto/pbgo/mocks/api_mockgen.pb.go +++ b/pkg/proto/pbgo/mocks/api_mockgen.pb.go @@ -1,5 +1,5 @@ // Code generated by MockGen. DO NOT EDIT. -// Source: /home/arthur/go/src/github.com/DataDog/datadog-agent/pkg/proto/pbgo/api.pb.go +// Source: /Users/arthur.bellal/go/src/github.com/DataDog/datadog-agent/pkg/proto/pbgo/api.pb.go // Package mock_pbgo is a generated GoMock package. package mock_pbgo diff --git a/tasks/go.py b/tasks/go.py index d928cbe377bab..c3ebcecb041de 100644 --- a/tasks/go.py +++ b/tasks/go.py @@ -437,6 +437,9 @@ def generate_protobuf(ctx): ctx.run(f"mockgen -source={pbgo_dir}/api.pb.go -destination={mockgen_out}/api_mockgen.pb.go") + # generate messagepack marshallers + ctx.run("msgp -file pkg/proto/msgpgo/key.go -o=pkg/proto/msgpgo/key_gen.go") + @task def reset(ctx):