Allow getting and patching wg pubkey from/to node repo.

+ Add missing 'trustStore' to NodeAttributes.hashCode

5 files changed, 68 insertions, 18 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeAttributes.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeAttributes.java
index 8c58022b67f..5d87c5dd3fc 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeAttributes.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeAttributes.java
@@ -4,6 +4,7 @@ package com.yahoo.vespa.hosted.node.admin.configserver.noderepository;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.DockerImage;
+import com.yahoo.config.provision.WireguardKey;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.reports.BaseReport;
 
 import java.time.Instant;
@@ -32,6 +33,7 @@ public class NodeAttributes {
     private Optional<Version> currentOsVersion = Optional.empty();
     private Optional<Instant> currentFirmwareCheck = Optional.empty();
     private List<TrustStoreItem> trustStore = List.of();
+    private Optional<WireguardKey> wireguardPubkey = Optional.empty();
     /** The list of reports to patch. A null value is used to remove the report. */
     private Map<String, JsonNode> reports = new TreeMap<>();
 
@@ -76,6 +78,16 @@ public class NodeAttributes {
         return this;
     }
 
+    public NodeAttributes withTrustStore(List<TrustStoreItem> trustStore) {
+        this.trustStore = List.copyOf(trustStore);
+        return this;
+    }
+
+    public NodeAttributes withWireguardPubkey(WireguardKey wireguardPubkey) {
+        this.wireguardPubkey = Optional.of(wireguardPubkey);
+        return this;
+    }
+
     public NodeAttributes withReports(Map<String, JsonNode> nodeReports) {
         this.reports = new TreeMap<>(nodeReports);
         return this;
@@ -119,6 +131,12 @@ public class NodeAttributes {
         return currentFirmwareCheck;
     }
 
+    public List<TrustStoreItem> getTrustStore() {
+        return trustStore;
+    }
+
+    public Optional<WireguardKey> getWireguardPubkey() { return wireguardPubkey; }
+
     public Map<String, JsonNode> getReports() {
         return reports;
     }
@@ -130,7 +148,7 @@ public class NodeAttributes {
     @Override
     public int hashCode() {
         return Objects.hash(hostId, restartGeneration, rebootGeneration, dockerImage, vespaVersion, currentOsVersion,
-                            currentFirmwareCheck, reports);
+                            currentFirmwareCheck, trustStore, wireguardPubkey, reports);
     }
 
     public boolean isEmpty() {
@@ -150,17 +168,9 @@ public class NodeAttributes {
                 && Objects.equals(vespaVersion, other.vespaVersion)
                 && Objects.equals(currentOsVersion, other.currentOsVersion)
                 && Objects.equals(currentFirmwareCheck, other.currentFirmwareCheck)
-                && Objects.equals(reports, other.reports)
-                && Objects.equals(trustStore, other.trustStore);
-    }
-
-    public NodeAttributes withTrustStore(List<TrustStoreItem> trustStore) {
-        this.trustStore = List.copyOf(trustStore);
-        return this;
-    }
-
-    public List<TrustStoreItem> getTrustStore() {
-        return trustStore;
+                && Objects.equals(trustStore, other.trustStore)
+                && Objects.equals(wireguardPubkey, other.wireguardPubkey)
+                && Objects.equals(reports, other.reports);
     }
 
     @Override
@@ -172,8 +182,9 @@ public class NodeAttributes {
                          vespaVersion.map(ver -> "vespaVersion=" + ver.toFullString()),
                          currentOsVersion.map(ver -> "currentOsVersion=" + ver.toFullString()),
                          currentFirmwareCheck.map(at -> "currentFirmwareCheck=" + at),
-                         Optional.ofNullable(reports.isEmpty() ? null : "reports=" + reports),
-                         Optional.ofNullable(trustStore.isEmpty() ? null : "trustStore=" + trustStore))
+                         Optional.ofNullable(trustStore.isEmpty() ? null : "trustStore=" + trustStore),
+                         Optional.ofNullable(wireguardPubkey.isEmpty() ? null : "wireguardPubkey=" + wireguardPubkey),
+                         Optional.ofNullable(reports.isEmpty() ? null : "reports=" + reports))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .collect(Collectors.joining(", ", "{", "}"));
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
index 1d4d8fa5281..e54b91d9c67 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.config.provision.WireguardKey;
 import com.yahoo.vespa.hosted.node.admin.task.util.file.DiskSize;
 
 import java.net.URI;
@@ -70,6 +71,8 @@ public class NodeSpec {
 
     private final List<TrustStoreItem> trustStore;
 
+    private final Optional<WireguardKey> wireguardPubkey;
+
     private final boolean wantToRebuild;
 
     public NodeSpec(
@@ -104,6 +107,7 @@ public class NodeSpec {
             Optional<URI> archiveUri,
             Optional<ApplicationId> exclusiveTo,
             List<TrustStoreItem> trustStore,
+            Optional<WireguardKey> wireguardPubkey,
             boolean wantToRebuild) {
 
         if (state == NodeState.active) {
@@ -146,6 +150,7 @@ public class NodeSpec {
         this.archiveUri = Objects.requireNonNull(archiveUri);
         this.exclusiveTo = Objects.requireNonNull(exclusiveTo);
         this.trustStore = Objects.requireNonNull(trustStore);
+        this.wireguardPubkey = Objects.requireNonNull(wireguardPubkey);
         this.wantToRebuild = wantToRebuild;
     }
 
@@ -296,6 +301,8 @@ public class NodeSpec {
         return trustStore;
     }
 
+    public Optional<WireguardKey> wireguardPubkey() { return wireguardPubkey; }
+
     public boolean wantToRebuild() {
         return wantToRebuild;
     }
@@ -336,6 +343,7 @@ public class NodeSpec {
                 Objects.equals(archiveUri, that.archiveUri) &&
                 Objects.equals(exclusiveTo, that.exclusiveTo) &&
                 Objects.equals(trustStore, that.trustStore) &&
+                Objects.equals(wireguardPubkey, that.wireguardPubkey) &&
                 Objects.equals(wantToRebuild, that.wantToRebuild);
     }
 
@@ -373,6 +381,7 @@ public class NodeSpec {
                 archiveUri,
                 exclusiveTo,
                 trustStore,
+                wireguardPubkey,
                 wantToRebuild);
     }
 
@@ -410,6 +419,7 @@ public class NodeSpec {
                 + " archiveUri=" + archiveUri
                 + " exclusiveTo=" + exclusiveTo
                 + " trustStore=" + trustStore
+                + " wireguardPubkey=" + wireguardPubkey
                 + " wantToRebuild=" + wantToRebuild
                 + " }";
     }
@@ -446,6 +456,7 @@ public class NodeSpec {
         private Optional<URI> archiveUri = Optional.empty();
         private Optional<ApplicationId> exclusiveTo = Optional.empty();
         private List<TrustStoreItem> trustStore = List.of();
+        private Optional<WireguardKey> wireguardPubkey = Optional.empty();
         private boolean wantToRebuild = false;
 
         public Builder() {}
@@ -481,6 +492,7 @@ public class NodeSpec {
             node.archiveUri.ifPresent(this::archiveUri);
             node.exclusiveTo.ifPresent(this::exclusiveTo);
             trustStore(node.trustStore);
+            node.wireguardPubkey.ifPresent(this::wireguardPubkey);
             wantToRebuild(node.wantToRebuild);
         }
 
@@ -664,6 +676,11 @@ public class NodeSpec {
             return this;
         }
 
+        public Builder wireguardPubkey(WireguardKey wireguardKey) {
+            wireguardPubkey = Optional.of(wireguardKey);
+            return this;
+        }
+
         public Builder wantToRebuild(boolean wantToRebuild) {
             this.wantToRebuild = wantToRebuild;
             return this;
@@ -677,6 +694,7 @@ public class NodeSpec {
             attributes.getRestartGeneration().ifPresent(this::currentRestartGeneration);
             // Always replace entire trust store
             trustStore(attributes.getTrustStore());
+            attributes.getWireguardPubkey().ifPresent(this::wireguardPubkey);
             this.reports.updateFromRawMap(attributes.getReports());
 
             return this;
@@ -790,7 +808,8 @@ public class NodeSpec {
                                 wantedRebootGeneration, currentRebootGeneration,
                                 wantedFirmwareCheck, currentFirmwareCheck, modelName,
                                 resources, realResources, ipAddresses, additionalIpAddresses,
-                                reports, events, parentHostname, archiveUri, exclusiveTo, trustStore, wantToRebuild);
+                                reports, events, parentHostname, archiveUri, exclusiveTo, trustStore,
+                                wireguardPubkey, wantToRebuild);
         }
 
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
index c15998a48df..cc41825fd5b 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepository.java
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.config.provision.WireguardKey;
 import com.yahoo.config.provision.host.FlavorOverrides;
 import com.yahoo.vespa.hosted.node.admin.configserver.ConfigServerApi;
 import com.yahoo.vespa.hosted.node.admin.configserver.HttpException;
@@ -200,6 +201,7 @@ public class RealNodeRepository implements NodeRepository {
                 Optional.ofNullable(node.archiveUri).map(URI::create),
                 Optional.ofNullable(node.exclusiveTo).map(ApplicationId::fromSerializedForm),
                 trustStore,
+                Optional.ofNullable(node.wireguardPubkey).map(WireguardKey::from),
                 node.wantToRebuild);
     }
 
@@ -316,6 +318,7 @@ public class RealNodeRepository implements NodeRepository {
         node.trustStore = nodeAttributes.getTrustStore().stream()
                 .map(item -> new NodeRepositoryNode.TrustStoreItem(item.fingerprint(), item.expiry().toEpochMilli()))
                 .toList();
+        node.wireguardPubkey = nodeAttributes.getWireguardPubkey().map(WireguardKey::value).orElse(null);
         Map<String, JsonNode> reports = nodeAttributes.getReports();
         node.reports = reports == null || reports.isEmpty() ? null : new TreeMap<>(reports);
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
index 434acca767a..eb7ceab6021 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/bindings/NodeRepositoryNode.java
@@ -90,6 +90,9 @@ public class NodeRepositoryNode {
     @JsonProperty("trustStore")
     @JsonInclude(JsonInclude.Include.NON_EMPTY)
     public List<TrustStoreItem> trustStore;
+    @JsonProperty("wireguardPubkey")
+    @JsonInclude(JsonInclude.Include.NON_EMPTY)
+    public String wireguardPubkey;
 
     @JsonProperty("reports")
     public Map<String, JsonNode> reports = null;
@@ -133,6 +136,7 @@ public class NodeRepositoryNode {
                ", exclusiveTo='" + exclusiveTo + '\'' +
                ", history=" + history +
                ", trustStore=" + trustStore +
+               ", wireguardPubkey=" + wireguardPubkey +
                ", reports=" + reports +
                '}';
     }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepositoryTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepositoryTest.java
index 661d3e6073e..672698c2e37 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepositoryTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/RealNodeRepositoryTest.java
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.CloudAccount;
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.config.provision.WireguardKey;
 import com.yahoo.config.provision.host.FlavorOverrides;
 import com.yahoo.vespa.hosted.node.admin.configserver.ConfigServerApi;
 import com.yahoo.vespa.hosted.node.admin.configserver.ConfigServerApiImpl;
@@ -24,7 +25,10 @@ import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 
-import static org.junit.jupiter.api.Assertions.*;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 
 /**
  * Tests the NodeRepository class used for talking to the node repository. It uses a mock from the node repository
@@ -128,12 +132,21 @@ public class RealNodeRepositoryTest {
 
     @Test
     void testUpdateNodeAttributes() {
-        String hostname = "host4.yahoo.com";
+        var hostname = "host4.yahoo.com";
+        var dockerImage = "registry.example.com/repo/image-1:6.2.3";
+        var wireguardKey = WireguardKey.from("111122223333444455556666777788889999000042c=");
+
         nodeRepositoryApi.updateNodeAttributes(
                 hostname,
                 new NodeAttributes()
                         .withRestartGeneration(1)
-                        .withDockerImage(DockerImage.fromString("registry.example.com/repo/image-1:6.2.3")));
+                        .withDockerImage(DockerImage.fromString(dockerImage))
+                        .withWireguardPubkey(wireguardKey));
+
+        NodeSpec hostSpec = nodeRepositoryApi.getOptionalNode(hostname).orElseThrow();
+        assertEquals(1, hostSpec.currentRestartGeneration().orElseThrow());
+        assertEquals(dockerImage, hostSpec.currentDockerImage().orElseThrow().asString());
+        assertEquals(wireguardKey.value(), hostSpec.wireguardPubkey().orElseThrow().value());
     }
 
     @Test