session: add GetGroupID and GetSessionIDs methods

This commit adds new getters: `GetGroupID` and `GetSessionIDs` to the
session store which can be used to query the newly added indexes to get
the associated group ID for a session ID or the associated set of
session IDs for a group ID.

Author: ellemouton

session/interface.go

@@ -128,6 +128,16 @@ func NewSession(id ID, localPrivKey *btcec.PrivateKey, label string, typ Type,
 	return sess, nil
 }
 
+// IDToGroupIndex defines an interface for the session ID to group ID index.
+type IDToGroupIndex interface {
+	// GetGroupID will return the group ID for the given session ID.
+	GetGroupID(sessionID ID) (ID, error)
+
+	// GetSessionIDs will return the set of session IDs that are in the
+	// group with the given ID.
+	GetSessionIDs(groupID ID) ([]ID, error)
+}
+
 // Store is the interface a persistent storage must implement for storing and
 // retrieving Terminal Connect sessions.
 type Store interface {
@@ -160,4 +170,6 @@ type Store interface {
 
 	// GetSessionByID fetches the session with the given ID.
 	GetSessionByID(id ID) (*Session, error)
+
+	IDToGroupIndex
 }

session/store.go

@@ -341,6 +341,91 @@ func (db *DB) GetUnusedIDAndKeyPair() (ID, *btcec.PrivateKey, error) {
 	return id, privKey, nil
 }
 
+// GetGroupID will return the group ID for the given session ID.
+//
+// NOTE: this is part of the IDToGroupIndex interface.
+func (db *DB) GetGroupID(sessionID ID) (ID, error) {
+	var groupID ID
+	err := db.View(func(tx *bbolt.Tx) error {
+		sessionBkt, err := getBucket(tx, sessionBucketKey)
+		if err != nil {
+			return err
+		}
+
+		idIndex := sessionBkt.Bucket(idIndexKey)
+		if idIndex == nil {
+			return ErrDBInitErr
+		}
+
+		sessionIDBkt := idIndex.Bucket(sessionID[:])
+		if sessionIDBkt == nil {
+			return fmt.Errorf("no index entry for session ID: %x",
+				sessionID)
+		}
+
+		groupIDBytes := sessionIDBkt.Get(groupIDKey)
+		if len(groupIDBytes) == 0 {
+			return fmt.Errorf("group ID not found for session "+
+				"ID %x", sessionID)
+		}
+
+		copy(groupID[:], groupIDBytes)
+
+		return nil
+	})
+	if err != nil {
+		return groupID, err
+	}
+
+	return groupID, nil
+}
+
+// GetSessionIDs will return the set of session IDs that are in the
+// group with the given ID.
+//
+// NOTE: this is part of the IDToGroupIndex interface.
+func (db *DB) GetSessionIDs(groupID ID) ([]ID, error) {
+	var sessionIDs []ID
+	err := db.View(func(tx *bbolt.Tx) error {
+		sessionBkt, err := getBucket(tx, sessionBucketKey)
+		if err != nil {
+			return err
+		}
+
+		groupIndexBkt := sessionBkt.Bucket(groupIDIndexKey)
+		if groupIndexBkt == nil {
+			return ErrDBInitErr
+		}
+
+		groupIDBkt := groupIndexBkt.Bucket(groupID[:])
+		if groupIDBkt == nil {
+			return fmt.Errorf("no sessions for group ID %v",
+				groupID)
+		}
+
+		sessionIDsBkt := groupIDBkt.Bucket(sessionIDKey)
+		if sessionIDsBkt == nil {
+			return fmt.Errorf("no sessions for group ID %v",
+				groupID)
+		}
+
+		return sessionIDsBkt.ForEach(func(_,
+			sessionIDBytes []byte) error {
+
+			var sessionID ID
+			copy(sessionID[:], sessionIDBytes)
+			sessionIDs = append(sessionIDs, sessionID)
+
+			return nil
+		})
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return sessionIDs, nil
+}
+
 // addIdToKeyPair inserts the mapping from session ID to session key into the
 // id-index bucket. An error is returned if an entry for this ID already exists.
 func addIDToKeyPair(sessionBkt *bbolt.Bucket, id ID, sessionKey []byte) error {

session/store_test.go

@@ -110,6 +110,68 @@ func TestLinkingSessions(t *testing.T) {
 	require.NoError(t, db.CreateSession(s2))
 }
 
+// TestIDToGroupIDIndex tests that the session-ID-to-group-ID and
+// group-ID-to-session-ID indexes work as expected by asserting the behaviour
+// of the GetGroupID and GetSessionIDs methods.
+func TestLinkedSessions(t *testing.T) {
+	// Set up a new DB.
+	db, err := NewDB(t.TempDir(), "test.db")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = db.Close()
+	})
+
+	// Create a few sessions. The first one is a new session and the two
+	// after are all linked to the prior one. All these sessions belong to
+	// the same group. The group ID is equivalent to the session ID of the
+	// first session.
+	s1 := newSession(t, db, "session 1", nil)
+	s2 := newSession(t, db, "session 2", &s1.GroupID)
+	s3 := newSession(t, db, "session 3", &s2.GroupID)
+
+	// Persist the sessions.
+	require.NoError(t, db.CreateSession(s1))
+	require.NoError(t, db.CreateSession(s2))
+	require.NoError(t, db.CreateSession(s3))
+
+	// Assert that the session ID to group ID index works as expected.
+	for _, s := range []*Session{s1, s2, s3} {
+		groupID, err := db.GetGroupID(s.ID)
+		require.NoError(t, err)
+		require.Equal(t, s1.ID, groupID)
+		require.Equal(t, s.GroupID, groupID)
+	}
+
+	// Assert that the group ID to session ID index works as expected.
+	sIDs, err := db.GetSessionIDs(s1.GroupID)
+	require.NoError(t, err)
+	require.EqualValues(t, []ID{s1.ID, s2.ID, s3.ID}, sIDs)
+
+	// To ensure that different groups don't interfere with each other,
+	// let's add another set of linked sessions not linked to the first.
+	s4 := newSession(t, db, "session 4", nil)
+	s5 := newSession(t, db, "session 5", &s4.GroupID)
+
+	require.NotEqual(t, s4.GroupID, s1.GroupID)
+
+	// Persist the sessions.
+	require.NoError(t, db.CreateSession(s4))
+	require.NoError(t, db.CreateSession(s5))
+
+	// Assert that the session ID to group ID index works as expected.
+	for _, s := range []*Session{s4, s5} {
+		groupID, err := db.GetGroupID(s.ID)
+		require.NoError(t, err)
+		require.Equal(t, s4.ID, groupID)
+		require.Equal(t, s.GroupID, groupID)
+	}
+
+	// Assert that the group ID to session ID index works as expected.
+	sIDs, err = db.GetSessionIDs(s5.GroupID)
+	require.NoError(t, err)
+	require.EqualValues(t, []ID{s4.ID, s5.ID}, sIDs)
+}
+
 func newSession(t *testing.T, db Store, label string,
 	linkedGroupID *ID) *Session {