1 files changed, 44 insertions, 6 deletions
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/IdentityDocumentSigner.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/IdentityDocumentSigner.java
index 019f73fc6bf..11b30585933 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/IdentityDocumentSigner.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/IdentityDocumentSigner.java
@@ -4,7 +4,10 @@ package com.yahoo.vespa.athenz.identityprovider.client;
 import com.yahoo.security.SignatureUtils;
 import com.yahoo.vespa.athenz.api.AthenzIdentity;
 import com.yahoo.vespa.athenz.api.AthenzService;
+import com.yahoo.vespa.athenz.identityprovider.api.DefaultSignedIdentityDocument;
+import com.yahoo.vespa.athenz.identityprovider.api.IdentityDocument;
 import com.yahoo.vespa.athenz.identityprovider.api.IdentityType;
+import com.yahoo.vespa.athenz.identityprovider.api.LegacySignedIdentityDocument;
 import com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument;
 import com.yahoo.vespa.athenz.identityprovider.api.VespaUniqueInstanceId;
 
@@ -19,7 +22,7 @@ import java.util.Base64;
 import java.util.Set;
 import java.util.TreeSet;
 
-import static com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument.DEFAULT_DOCUMENT_VERSION;
+import static com.yahoo.vespa.athenz.identityprovider.api.SignedIdentityDocument.LEGACY_DEFAULT_DOCUMENT_VERSION;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 /**
@@ -29,8 +32,25 @@ import static java.nio.charset.StandardCharsets.UTF_8;
  */
 public class IdentityDocumentSigner {
 
+    public String generateSignature(String identityDocumentData, PrivateKey privateKey) {
+        try {
+            Signature signer = SignatureUtils.createSigner(privateKey);
+            signer.initSign(privateKey);
+            signer.update(identityDocumentData.getBytes(UTF_8));
+            byte[] signature = signer.sign();
+            return Base64.getEncoder().encodeToString(signature);
+        } catch (GeneralSecurityException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public String generateLegacySignature(IdentityDocument doc, PrivateKey privateKey) {
+        return generateSignature(doc.providerUniqueId(), doc.providerService(), doc.configServerHostname(),
+                                 doc.instanceHostname(), doc.createdAt(), doc.ipAddresses(), doc.identityType(), privateKey, doc.serviceIdentity());
+    }
+
     // Cluster type is ignored due to old Vespa versions not forwarding unknown fields in signed identity document
-    public String generateSignature(VespaUniqueInstanceId providerUniqueId,
+    private String generateSignature(VespaUniqueInstanceId providerUniqueId,
                                     AthenzService providerService,
                                     String configServerHostname,
                                     String instanceHostname,
@@ -54,14 +74,32 @@ public class IdentityDocumentSigner {
     }
 
     public boolean hasValidSignature(SignedIdentityDocument doc, PublicKey publicKey) {
+        if (doc instanceof LegacySignedIdentityDocument signedDoc) {
+            return validateLegacySignature(signedDoc, publicKey);
+        } else if (doc instanceof DefaultSignedIdentityDocument signedDoc) {
+            try {
+                Signature signer = SignatureUtils.createVerifier(publicKey);
+                signer.initVerify(publicKey);
+                signer.update(signedDoc.data().getBytes(UTF_8));
+                return signer.verify(Base64.getDecoder().decode(doc.signature()));
+            } catch (GeneralSecurityException e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            throw new IllegalArgumentException("Unknown identity document type: " + doc.getClass().getName());
+        }
+    }
+
+    private boolean validateLegacySignature(SignedIdentityDocument doc, PublicKey publicKey) {
         try {
+            IdentityDocument iddoc = doc.identityDocument();
             Signature signer = SignatureUtils.createVerifier(publicKey);
             signer.initVerify(publicKey);
             writeToSigner(
-                    signer, doc.providerUniqueId(), doc.providerService(), doc.configServerHostname(),
-                    doc.instanceHostname(), doc.createdAt(), doc.ipAddresses(), doc.identityType());
-            if (doc.documentVersion() >= DEFAULT_DOCUMENT_VERSION) {
-                writeToSigner(signer, doc.serviceIdentity());
+                    signer, iddoc.providerUniqueId(), iddoc.providerService(), iddoc.configServerHostname(),
+                    iddoc.instanceHostname(), iddoc.createdAt(), iddoc.ipAddresses(), iddoc.identityType());
+            if (doc.documentVersion() >= LEGACY_DEFAULT_DOCUMENT_VERSION) {
+                writeToSigner(signer, iddoc.serviceIdentity());
             }
             return signer.verify(Base64.getDecoder().decode(doc.signature()));
         } catch (GeneralSecurityException e) {