From da528f2d9b53286c2e162345945bf4a43f09ec26 Mon Sep 17 00:00:00 2001
From: Leander Beernaert <leander.beernaert@proton.ch>
Date: Tue, 25 Jul 2023 09:59:26 +0200
Subject: [PATCH] feat(GODT-2801): Add Key Decryption helpers to
 `identityservice.State`

---
 internal/services/useridentity/state.go | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/internal/services/useridentity/state.go b/internal/services/useridentity/state.go
index ac8e90d9..6d5bbdb4 100644
--- a/internal/services/useridentity/state.go
+++ b/internal/services/useridentity/state.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"crypto/subtle"
 	"fmt"
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"strings"
 
 	"github.com/ProtonMail/go-proton-api"
@@ -243,3 +244,16 @@ func (s *State) CheckAuth(email string, password []byte, bridgePassProvider Brid
 
 	return "", fmt.Errorf("invalid email")
 }
+
+func (s *State) WithAddrKR(addrID string, keyPass []byte, fn func(userKR, addrKR *crypto.KeyRing) error) error {
+	addr, ok := s.Addresses[addrID]
+	if !ok {
+		return fmt.Errorf("address not found")
+	}
+
+	return usertypes.WithAddrKR(s.User, addr, keyPass, fn)
+}
+
+func (s *State) WithAddrKRs(keyPass []byte, fn func(*crypto.KeyRing, map[string]*crypto.KeyRing) error) error {
+	return usertypes.WithAddrKRs(s.User, s.Addresses, keyPass, fn)
+}