firewall: map session ID to group ID in privacy mapper interceptor

firewall/privacy_mapper.go

@@ -46,16 +46,22 @@ var _ mid.RequestInterceptor = (*PrivacyMapper)(nil)
 // PrivacyMapper is a RequestInterceptor that maps any pseudo names in certain
 // requests to their real values and vice versa for responses.
 type PrivacyMapper struct {
-	newDB    firewalldb.NewPrivacyMapDB
-	randIntn func(int) (int, error)
+	newDB            firewalldb.NewPrivacyMapDB
+	randIntn         func(int) (int, error)
+	sessionIDIndexDB session.IDToGroupIndex
 }
 
 // NewPrivacyMapper returns a new instance of PrivacyMapper. The randIntn
 // function is used to draw randomness for request field obfuscation.
 func NewPrivacyMapper(newDB firewalldb.NewPrivacyMapDB,
-	randIntn func(int) (int, error)) *PrivacyMapper {
+	randIntn func(int) (int, error),
+	sessionIDIndexDB session.IDToGroupIndex) *PrivacyMapper {
 
-	return &PrivacyMapper{newDB: newDB, randIntn: randIntn}
+	return &PrivacyMapper{
+		newDB:            newDB,
+		randIntn:         randIntn,
+		sessionIDIndexDB: sessionIDIndexDB,
+	}
 }
 
 // Name returns the name of the interceptor.
@@ -91,6 +97,12 @@ func (p *PrivacyMapper) Intercept(ctx context.Context,
 		return nil, fmt.Errorf("could not extract ID from macaroon")
 	}
 
+	// Get group ID for session ID.
+	groupID, err := p.sessionIDIndexDB.GetGroupID(sessionID)
+	if err != nil {
+		return nil, err
+	}
+
 	log.Tracef("PrivacyMapper: Intercepting %v", ri)
 
 	switch r := req.InterceptType.(type) {
@@ -108,7 +120,7 @@ func (p *PrivacyMapper) Intercept(ctx context.Context,
 		}
 
 		replacement, err := p.checkAndReplaceIncomingRequest(
-			ctx, r.Request.MethodFullUri, msg, sessionID,
+			ctx, r.Request.MethodFullUri, msg, groupID,
 		)
 		if err != nil {
 			return mid.RPCErr(req, err)
@@ -142,7 +154,7 @@ func (p *PrivacyMapper) Intercept(ctx context.Context,
 		}
 
 		replacement, err := p.replaceOutgoingResponse(
-			ctx, r.Response.MethodFullUri, msg, sessionID,
+			ctx, r.Response.MethodFullUri, msg, groupID,
 		)
 		if err != nil {
 			return mid.RPCErr(req, err)
@@ -167,10 +179,10 @@ func (p *PrivacyMapper) Intercept(ctx context.Context,
 // checkAndReplaceIncomingRequest inspects an incoming request and optionally
 // modifies some of the request parameters.
 func (p *PrivacyMapper) checkAndReplaceIncomingRequest(ctx context.Context,
-	uri string, req proto.Message, sessionID session.ID) (proto.Message,
+	uri string, req proto.Message, groupID session.ID) (proto.Message,
 	error) {
 
-	db := p.newDB(sessionID)
+	db := p.newDB(groupID)
 
 	// If we don't have a handler for the URI, we don't allow the request
 	// to go through.
@@ -193,9 +205,9 @@ func (p *PrivacyMapper) checkAndReplaceIncomingRequest(ctx context.Context,
 // replaceOutgoingResponse inspects the responses before sending them out to the
 // client and replaces them if needed.
 func (p *PrivacyMapper) replaceOutgoingResponse(ctx context.Context, uri string,
-	resp proto.Message, sessionID session.ID) (proto.Message, error) {
+	resp proto.Message, groupID session.ID) (proto.Message, error) {
 
-	db := p.newDB(sessionID)
+	db := p.newDB(groupID)
 
 	// If we don't have a handler for the URI, we don't allow the response
 	// to go to avoid accidental leaks.

firewall/privacy_mapper_test.go

@@ -2,6 +2,7 @@ package firewall
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -292,9 +293,12 @@ func TestPrivacyMapper(t *testing.T) {
 
 	db := newMockDB(t, mapPreloadRealToPseudo, sessionID)
 
+	err = db.AddSessionAndGroupIDPair(sessionID, sessionID)
+	require.NoError(t, err)
+
 	// randIntn is used for deterministic testing.
 	randIntn := func(n int) (int, error) { return 100, nil }
-	p := NewPrivacyMapper(db.NewSessionDB, randIntn)
+	p := NewPrivacyMapper(db.NewSessionDB, randIntn, db)
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
@@ -355,7 +359,7 @@ func TestPrivacyMapper(t *testing.T) {
 		rawMsg, err := proto.Marshal(msg)
 		require.NoError(t, err)
 
-		p = NewPrivacyMapper(db.NewSessionDB, CryptoRandIntn)
+		p = NewPrivacyMapper(db.NewSessionDB, CryptoRandIntn, db)
 		require.NoError(t, err)
 
 		// We test the independent outgoing amount (incoming amount
@@ -440,12 +444,21 @@ func TestPrivacyMapper(t *testing.T) {
 	})
 }
 
-type mockDB map[string]*mockPrivacyMapDB
+type mockDB struct {
+	privDB map[string]*mockPrivacyMapDB
+
+	sessionIDIndex map[session.ID]session.ID
+	groupIDIndex   map[session.ID][]session.ID
+}
 
 func newMockDB(t *testing.T, preloadRealToPseudo map[string]string,
 	sessID session.ID) mockDB {
 
-	db := make(mockDB)
+	db := mockDB{
+		privDB:         make(map[string]*mockPrivacyMapDB),
+		sessionIDIndex: make(map[session.ID]session.ID),
+		groupIDIndex:   make(map[session.ID][]session.ID),
+	}
 	sessDB := db.NewSessionDB(sessID)
 
 	_ = sessDB.Update(func(tx firewalldb.PrivacyMapTx) error {
@@ -459,17 +472,41 @@ func newMockDB(t *testing.T, preloadRealToPseudo map[string]string,
 }
 
 func (m mockDB) NewSessionDB(sessionID session.ID) firewalldb.PrivacyMapDB {
-	db, ok := m[string(sessionID[:])]
+	db, ok := m.privDB[string(sessionID[:])]
 	if ok {
 		return db
 	}
 
 	newDB := newMockPrivacyMapDB()
-	m[string(sessionID[:])] = newDB
+	m.privDB[string(sessionID[:])] = newDB
 
 	return newDB
 }
 
+func (m mockDB) AddSessionAndGroupIDPair(sessionID, groupID session.ID) error {
+	m.sessionIDIndex[sessionID] = groupID
+	m.groupIDIndex[groupID] = append(m.groupIDIndex[groupID], sessionID)
+	return nil
+}
+
+func (m mockDB) GetGroupID(sessionID session.ID) (session.ID, error) {
+	groupID, ok := m.sessionIDIndex[sessionID]
+	if !ok {
+		return session.ID{}, fmt.Errorf("group ID not found")
+	}
+
+	return groupID, nil
+}
+
+func (m mockDB) GetSessionIDs(groupID session.ID) ([]session.ID, error) {
+	sessionIDs, ok := m.groupIDIndex[groupID]
+	if !ok {
+		return nil, fmt.Errorf("group ID not found")
+	}
+
+	return sessionIDs, nil
+}
+
 func newMockPrivacyMapDB() *mockPrivacyMapDB {
 	return &mockPrivacyMapDB{
 		r2p: make(map[string]string),