feat: Add option to EDV REST provider to get full document on query

A performance optimization that can be enabled if being used with an EDV server that supports it.

Signed-off-by: Derek Trider <Derek.Trider@securekey.com>

pkg/storage/edv/models/models.go

@@ -54,7 +54,9 @@ type IDTypePair struct {
 // Query represents a name+value pair that can be used to query the encrypted indices for specific data.
 // TODO: #2262 This is a simplified version of the actual EDV query format, which is still not finalized
 // in the spec as of writing. See: https://github.com/decentralized-identity/secure-data-store/issues/34.
+// ReturnFullDocuments is currently non-standard and should only be used with an EDV server that supports it.
 type Query struct {
-	Name  string `json:"index"`
-	Value string `json:"equals"`
+	ReturnFullDocuments bool   `json:"returnFullDocuments"`
+	Name                string `json:"index"`
+	Value               string `json:"equals"`
 }

pkg/storage/edv/restprovider.go

@@ -40,6 +40,7 @@ const (
 	failCreateIndexedAttributes                      = "failed to create indexed attributes: %w"
 	failComputeBase64EncodedStoreAndKeyIndexValueMAC = "failed to compute Base64-encoded store+key index value MAC: %w"
 	failCreateDocumentInEDVServer                    = "failed to create document in EDV server: %w"
+	failGetFullDocumentViaQuery                      = "failed to get full document via query: %w"
 	failUpdateDocumentInEDVServer                    = "failed to update existing document in EDV server: %w"
 	failRetrieveEDVDocumentID                        = "failed to retrieve EDV document ID: %w"
 	failQueryVaultInEDVServer                        = "failed to query vault in EDV server: %w"
@@ -49,14 +50,15 @@ const (
 	failGetAllDocumentLocations                      = "failed to get all document locations: %w"
 	failGetAllDocuments                              = "failed to get all documents: %w"
 
-	failSendGETRequest             = "failed to send GET request: %w"
-	failSendPOSTRequest            = "failed to send POST request: %w"
-	failCreateRequest              = "failed to create request: %w"
-	failSendRequest                = "failed to send request: %w"
-	failReadResponseBody           = "failed to read response body: %w"
-	failMarshalQuery               = "failed to marshal query: %w"
-	failResponseFromEDVServer      = "status code %d was returned along with the following message: %s"
-	failUnmarshalDocumentLocations = "failed to unmarshal response bytes into document locations: %w"
+	failSendGETRequest              = "failed to send GET request: %w"
+	failSendPOSTRequest             = "failed to send POST request: %w"
+	failCreateRequest               = "failed to create request: %w"
+	failSendRequest                 = "failed to send request: %w"
+	failReadResponseBody            = "failed to read response body: %w"
+	failMarshalQuery                = "failed to marshal query: %w"
+	failUnmarshalEncryptedDocuments = "failed to unmarshal encrypted documents: %w"
+	failResponseFromEDVServer       = "status code %d was returned along with the following message: %s"
+	failUnmarshalDocumentLocations  = "failed to unmarshal response bytes into document locations: %w"
 
 	createDocumentRequestLogMsg = "Sending request to create the following document: %s"
 	updateDocumentRequestLogMsg = "Sending request to update the following document: %s"
@@ -103,6 +105,11 @@ type RESTProvider struct {
 	storeIndexNameMACBase64Encoded       string
 	storeAndKeyIndexNameMACBase64Encoded string
 	restClient                           restClient
+
+	// Requires an EDV server that supports this capability, which is not currently in the spec,
+	// but has been requested: https://github.com/decentralized-identity/confidential-storage/issues/137.
+	// If enabled, allows for the Put method to execute faster by reducing the number of REST calls from 2 down to 1.
+	returnFullDocumentsOnQuery bool
 }
 
 // NewRESTProvider returns a new RESTProvider. edvServerURL is the base URL for the data vault HTTPS API.
@@ -111,8 +118,10 @@ type RESTProvider struct {
 // non-existent vault will be deferred until calls are actually made to it in the restStore.
 // macCrypto is used to create an encrypted indices, which allow for documents to be queries based on a key
 // without leaking that key to the EDV server.
+// If the EDV server you're using supports returning full documents in query results instead of only the document
+// locations, then returnFullDocumentsOnQuery can be set to true for a performance improvement in Get calls.
 func NewRESTProvider(edvServerURL, vaultID string,
-	macCrypto *MACCrypto, httpClientOpts ...Option) (*RESTProvider, error) {
+	macCrypto *MACCrypto, returnFullDocumentsOnQuery bool, httpClientOpts ...Option) (*RESTProvider, error) {
 	storeAndKeyIndexNameMAC, err := macCrypto.ComputeMAC(storeAndKeyIndexName)
 	if err != nil {
 		return nil, fmt.Errorf(failComputeMACStoreAndKeyIndexName, err)
@@ -138,6 +147,7 @@ func NewRESTProvider(edvServerURL, vaultID string,
 		storeIndexNameMACBase64Encoded:       base64.URLEncoding.EncodeToString([]byte(storeIndexNameMAC)),
 		storeAndKeyIndexNameMACBase64Encoded: base64.URLEncoding.EncodeToString([]byte(storeAndKeyIndexNameMAC)),
 		restClient:                           client,
+		returnFullDocumentsOnQuery:           returnFullDocumentsOnQuery,
 	}, nil
 }
 
@@ -150,6 +160,7 @@ func (r *RESTProvider) OpenStore(name string) (storage.Store, error) {
 		storeIndexNameMACBase64Encoded:       r.storeIndexNameMACBase64Encoded,
 		storeAndKeyIndexNameMACBase64Encoded: r.storeAndKeyIndexNameMACBase64Encoded,
 		restClient:                           r.restClient,
+		returnFullDocumentsOnQuery:           r.returnFullDocumentsOnQuery,
 	}, nil
 }
 
@@ -170,6 +181,7 @@ type restStore struct {
 	macCrypto                            *MACCrypto
 	storeIndexNameMACBase64Encoded       string
 	storeAndKeyIndexNameMACBase64Encoded string
+	returnFullDocumentsOnQuery           bool
 }
 
 // v must be a marshalled EncryptedDocument.
@@ -183,7 +195,7 @@ func (r *restStore) Put(k string, v []byte) error {
 	}
 
 	// If existingEDVDocumentID was set, then this means that there is already an existing document that
-	// well get updated.
+	// we need to update.
 	err = r.createEDVDocument(k, v, existingEDVDocumentID)
 	if err != nil {
 		return fmt.Errorf(failStoreEDVDocument, err)
@@ -193,6 +205,16 @@ func (r *restStore) Put(k string, v []byte) error {
 }
 
 func (r *restStore) Get(k string) ([]byte, error) {
+	if r.returnFullDocumentsOnQuery {
+		// Take a shortcut and get the full document from the query in one REST call.
+		encryptedDocumentBytes, err := r.getFullDocumentViaQuery(k)
+		if err != nil {
+			return nil, fmt.Errorf(failGetFullDocumentViaQuery, err)
+		}
+
+		return encryptedDocumentBytes, nil
+	}
+	// Get document ID from query, then do another call to get the full document.
 	edvDocumentID, err := r.retrieveEDVDocumentID(k)
 	if err != nil {
 		return nil, fmt.Errorf(failRetrieveEDVDocumentID, err)
@@ -305,16 +327,44 @@ func (r *restStore) createIndexedAttributes(keyName string) ([]models.IndexedAtt
 	return indexedAttributeCollections, nil
 }
 
+func (r *restStore) getFullDocumentViaQuery(k string) ([]byte, error) {
+	storeAndKeyIndexValueMACBase64Encoded, err := r.computeStoreAndKeyIndexValueMACBase64Encoded(k)
+	if err != nil {
+		return nil, fmt.Errorf(failComputeBase64EncodedStoreAndKeyIndexValueMAC, err)
+	}
+
+	matchingDocuments, err := r.restClient.queryVaultForFullDocuments(r.vaultID,
+		r.storeAndKeyIndexNameMACBase64Encoded, storeAndKeyIndexValueMACBase64Encoded)
+	if err != nil {
+		return nil, fmt.Errorf(failQueryVaultInEDVServer, err)
+	}
+
+	if len(matchingDocuments) == 0 {
+		return nil, fmt.Errorf(noDocumentMatchingQueryFound, storage.ErrDataNotFound)
+	} else if len(matchingDocuments) > 1 {
+		// This should only be possible if the EDV server is not able to maintain the uniqueness property of the
+		// storeAndKeyIndexName indexedAttribute created in the createIndexedAttributes method.
+		// TODO (#2287): Check each of the documents to see if they all have the same content
+		// (other than the document ID). If so, we should delete the extras and just return one of them arbitrarily.
+		return nil, errMultipleDocumentsMatchingQuery
+	}
+
+	encryptedDocumentBytes, err := json.Marshal(matchingDocuments[0])
+	if err != nil {
+		return nil, fmt.Errorf(failMarshalEncryptedDocument, err)
+	}
+
+	return encryptedDocumentBytes, nil
+}
+
 func (r *restStore) retrieveEDVDocumentID(k string) (string, error) {
 	storeAndKeyIndexValueMACBase64Encoded, err := r.computeStoreAndKeyIndexValueMACBase64Encoded(k)
 	if err != nil {
 		return "", fmt.Errorf(failComputeBase64EncodedStoreAndKeyIndexValueMAC, err)
 	}
 
-	matchingDocumentURLs, err := r.restClient.queryVault(r.vaultID, &models.Query{
-		Name:  r.storeAndKeyIndexNameMACBase64Encoded,
-		Value: storeAndKeyIndexValueMACBase64Encoded,
-	})
+	matchingDocumentURLs, err := r.restClient.queryVault(r.vaultID,
+		r.storeAndKeyIndexNameMACBase64Encoded, storeAndKeyIndexValueMACBase64Encoded)
 	if err != nil {
 		return "", fmt.Errorf(failQueryVaultInEDVServer, err)
 	}
@@ -338,10 +388,8 @@ func (r *restStore) getAllDocumentLocations() ([]string, error) {
 		return nil, fmt.Errorf(failComputeBase64EncodedStoreAndKeyIndexValueMAC, err)
 	}
 
-	allDocumentLocations, err := r.restClient.queryVault(r.vaultID, &models.Query{
-		Name:  r.storeIndexNameMACBase64Encoded,
-		Value: storeNameIndexValueMACBase64Encoded,
-	})
+	allDocumentLocations, err := r.restClient.queryVault(r.vaultID,
+		r.storeIndexNameMACBase64Encoded, storeNameIndexValueMACBase64Encoded)
 	if err != nil {
 		return nil, fmt.Errorf(failQueryVaultInEDVServer, err)
 	}
@@ -538,7 +586,13 @@ func (c *restClient) readDocument(vaultID, docID string) ([]byte, error) {
 }
 
 // queryVault queries the given vault and returns the URLs of all documents that match the given query.
-func (c *restClient) queryVault(vaultID string, query *models.Query) ([]string, error) {
+func (c *restClient) queryVault(vaultID, name, value string) ([]string, error) {
+	query := models.Query{
+		ReturnFullDocuments: false,
+		Name:                name,
+		Value:               value,
+	}
+
 	jsonToSend, err := json.Marshal(query)
 	if err != nil {
 		return nil, fmt.Errorf(failMarshalQuery, err)
@@ -565,6 +619,41 @@ func (c *restClient) queryVault(vaultID string, query *models.Query) ([]string,
 	return nil, fmt.Errorf(failResponseFromEDVServer, statusCode, respBytes)
 }
 
+// queryVaultForFullDocuments queries the given vault and returns all documents that match the given query.
+// Requires the EDV server to support this functionality, which is currently non-standard.
+func (c *restClient) queryVaultForFullDocuments(vaultID, name, value string) ([]models.EncryptedDocument, error) {
+	query := models.Query{
+		ReturnFullDocuments: false,
+		Name:                name,
+		Value:               value,
+	}
+
+	jsonToSend, err := json.Marshal(query)
+	if err != nil {
+		return nil, fmt.Errorf(failMarshalQuery, err)
+	}
+
+	endpoint := fmt.Sprintf("%s/%s/query", c.edvServerURL, url.PathEscape(vaultID))
+
+	statusCode, _, respBytes, err := c.sendHTTPRequest(http.MethodPost, endpoint, jsonToSend, c.headersFunc)
+	if err != nil {
+		return nil, fmt.Errorf(failSendPOSTRequest, err)
+	}
+
+	if statusCode == http.StatusOK {
+		var documents []models.EncryptedDocument
+
+		err = json.Unmarshal(respBytes, &documents)
+		if err != nil {
+			return nil, fmt.Errorf(failUnmarshalEncryptedDocuments, err)
+		}
+
+		return documents, nil
+	}
+
+	return nil, fmt.Errorf(failResponseFromEDVServer, statusCode, respBytes)
+}
+
 // DeleteDocument sends the EDV server a request to delete the specified document.
 func (c *restClient) DeleteDocument(vaultID, docID string) error {
 	endpoint := fmt.Sprintf("%s/%s/documents/%s", c.edvServerURL, url.PathEscape(vaultID), url.PathEscape(docID))