asserts: extend optional primary keys support to the in-memory backend

Merge pull request #11446 from pedronis/asserts-opt-primary-keys-2 As this one is in-process there are not backward compatibility concerns.

2 files changed, 112 insertions, 2 deletions

diff --git a/asserts/membackstore.go b/asserts/membackstore.go
index 991971c4b5..a23a72d357 100644
--- a/asserts/membackstore.go
+++ b/asserts/membackstore.go
@@ -1,7 +1,7 @@
 // -*- Mode: Go; indent-tabs-mode: t -*-
 
 /*
- * Copyright (C) 2016-2020 Canonical Ltd
+ * Copyright (C) 2016-2022 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
@@ -235,9 +235,23 @@ func (mbs *memoryBackstore) Get(assertType *AssertionType, key []string, maxForm
	mbs.mu.RLock()
	defer mbs.mu.RUnlock()
 
+	n := len(assertType.PrimaryKey)
+	if len(key) > n {
+	return nil, fmt.Errorf("internal error: Backstore.Get given a key longer than expected for %q: %v", assertType.Name, key)
+	}
+
	internalKey := make([]string, 1+len(assertType.PrimaryKey))
	internalKey[0] = assertType.Name
	copy(internalKey[1:], key)
+	if len(key) < n {
+	for kopt := len(key); kopt < n; kopt++ {
+	defl := assertType.OptionalPrimaryKeyDefaults[assertType.PrimaryKey[kopt]]
+	if defl == "" {
+	return nil, fmt.Errorf("internal error: Backstore.Get given a key missing mandatory elements for %q: %v", assertType.Name, key)
+	}
+	internalKey[kopt+1] = defl
+	}
+	}
 
	a, err := mbs.top.get(internalKey, maxFormat)
	if err == errNotFound {
diff --git a/asserts/membackstore_test.go b/asserts/membackstore_test.go
index ceced28d25..a5d7f2ed59 100644
--- a/asserts/membackstore_test.go
+++ b/asserts/membackstore_test.go
@@ -1,7 +1,7 @@
 // -*- Mode: Go; indent-tabs-mode: t -*-
 
 /*
- * Copyright (C) 2016-2020 Canonical Ltd
+ * Copyright (C) 2016-2022 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
@@ -20,6 +20,8 @@
 package asserts_test
 
 import (
+	"strings"
+
	. "gopkg.in/check.v1"
 
	"github.com/snapcore/snapd/asserts"
@@ -537,3 +539,97 @@ func (mbss *memBackstoreSuite) TestSequenceMemberAfter(c *C) {
	Type: asserts.TestOnlySeqType,
	})
 }
+
+func (mbss *memBackstoreSuite) TestOptionalPrimaryKeys(c *C) {
+	r := asserts.AddOptionalPrimaryKey(asserts.TestOnlyType, "opt1", "o1-defl")
+	defer r()
+	bs := mbss.bs
+
+	a1, err := asserts.Decode([]byte("type: test-only\n" +
+	"authority-id: auth-id1\n" +
+	"primary-key: k1\n" +
+	"marker: a1\n" +
+	"sign-key-sha3-384: Jv8_JiHiIzJVcO9M55pPdqSDWUvuhfDIBJUS-3VW7F_idjix7Ffn5qMxB21ZQuij" +
+	"\n\n" +
+	"AXNpZw=="))
+	c.Assert(err, IsNil)
+	err = bs.Put(asserts.TestOnlyType, a1)
+	c.Assert(err, IsNil)
+
+	a2, err := asserts.Decode([]byte("type: test-only\n" +
+	"authority-id: auth-id1\n" +
+	"primary-key: k2\n" +
+	"opt1: A\n" +
+	"marker: a2\n" +
+	"sign-key-sha3-384: Jv8_JiHiIzJVcO9M55pPdqSDWUvuhfDIBJUS-3VW7F_idjix7Ffn5qMxB21ZQuij" +
+	"\n\n" +
+	"AXNpZw=="))
+	c.Assert(err, IsNil)
+	err = bs.Put(asserts.TestOnlyType, a2)
+	c.Assert(err, IsNil)
+
+	a, err := bs.Get(asserts.TestOnlyType, []string{"k1"}, 0)
+	c.Assert(err, IsNil)
+	c.Check(a.Ref().PrimaryKey, DeepEquals, []string{"k1", "o1-defl"})
+	c.Check(a.HeaderString("marker"), Equals, "a1")
+
+	a, err = bs.Get(asserts.TestOnlyType, []string{"k1", "o1-defl"}, 0)
+	c.Assert(err, IsNil)
+	c.Check(a.Ref().PrimaryKey, DeepEquals, []string{"k1", "o1-defl"})
+	c.Check(a.HeaderString("marker"), Equals, "a1")
+
+	a, err = bs.Get(asserts.TestOnlyType, []string{"k2", "A"}, 0)
+	c.Assert(err, IsNil)
+	c.Check(a.Ref().PrimaryKey, DeepEquals, []string{"k2", "A"})
+	c.Check(a.HeaderString("marker"), Equals, "a2")
+
+	a, err = bs.Get(asserts.TestOnlyType, []string{"k2"}, 0)
+	c.Check(err, DeepEquals, &asserts.NotFoundError{
+	Type: asserts.TestOnlyType,
+	})
+
+	a, err = bs.Get(asserts.TestOnlyType, []string{}, 0)
+	c.Check(err, ErrorMatches, `internal error: Backstore.Get given a key missing mandatory elements for "test-only":.*`)
+
+	var found map[string]string
+	cb := func(a asserts.Assertion) {
+	if found == nil {
+	found = make(map[string]string)
+	}
+	found[strings.Join(a.Ref().PrimaryKey, "/")] = a.HeaderString("marker")
+
+	}
+	err = mbss.bs.Search(asserts.TestOnlyType, nil, cb, 0)
+	c.Assert(err, IsNil)
+	c.Check(found, DeepEquals, map[string]string{
+	"k1/o1-defl": "a1",
+	"k2/A": "a2",
+	})
+
+	found = nil
+	err = mbss.bs.Search(asserts.TestOnlyType, map[string]string{
+	"primary-key": "k1",
+	}, cb, 0)
+	c.Assert(err, IsNil)
+	c.Check(found, DeepEquals, map[string]string{
+	"k1/o1-defl": "a1",
+	})
+
+	found = nil
+	err = mbss.bs.Search(asserts.TestOnlyType, map[string]string{
+	"opt1": "o1-defl",
+	}, cb, 0)
+	c.Assert(err, IsNil)
+	c.Check(found, DeepEquals, map[string]string{
+	"k1/o1-defl": "a1",
+	})
+
+	found = nil
+	err = mbss.bs.Search(asserts.TestOnlyType, map[string]string{
+	"opt1": "A",
+	}, cb, 0)
+	c.Assert(err, IsNil)
+	c.Check(found, DeepEquals, map[string]string{
+	"k2/A": "a2",
+	})
+}