001/*
002 * Copyright 2022-2026 Revetware LLC.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 * http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package com.soklet;
018
019import org.jspecify.annotations.NonNull;
020
021import javax.annotation.concurrent.NotThreadSafe;
022import javax.annotation.concurrent.ThreadSafe;
023import java.util.Arrays;
024import java.util.Collection;
025import java.util.LinkedHashMap;
026import java.util.Map;
027import java.util.Set;
028
029import static java.util.Objects.requireNonNull;
030
031/**
032 * Immutable initial MCP request-state protection keyring.
033 * <p>
034 * Exactly one key is initially active for sealing. Other entries are initially
035 * verification-only. Building a server copies the complete ring, including
036 * key material, into independent server-owned live state. Runtime rotation is
037 * available only through that server's {@link McpProtectionKeyringManager}.
038 *
039 * @author <a href="https://www.revetkn.com">Mark Allen</a>
040 */
041@ThreadSafe
042public final class McpProtectionKeyring {
043        @NonNull
044        private final McpProtectionKey activeKey;
045        @NonNull
046        private final Map<@NonNull String, @NonNull McpProtectionKey> verificationKeys;
047
048        /**
049         * Vends a builder primed with the initial active key.
050         *
051         * @param activeKey initial active key
052         * @return keyring builder
053         */
054        @NonNull
055        public static Builder withActiveKey(@NonNull McpProtectionKey activeKey) {
056                return new Builder(activeKey);
057        }
058
059        private McpProtectionKeyring(@NonNull Builder builder) {
060                this.activeKey = copyOf(builder.activeKey);
061                LinkedHashMap<@NonNull String, @NonNull McpProtectionKey> copies =
062                                new LinkedHashMap<>();
063                builder.verificationKeys.forEach(
064                                (keyId, key) -> copies.put(keyId, copyOf(key)));
065                this.verificationKeys = Map.copyOf(copies);
066        }
067
068        /** @return non-secret active key ID */
069        @NonNull
070        public String getActiveKeyId() {
071                return this.activeKey.getKeyId();
072        }
073
074        /** @return immutable verification-only key ID set */
075        @NonNull
076        public Set<@NonNull String> getVerificationKeyIds() {
077                return this.verificationKeys.keySet();
078        }
079
080        @NonNull
081        McpProtectionKey initialActiveKey() {
082                return this.activeKey;
083        }
084
085        @NonNull
086        Collection<@NonNull McpProtectionKey> initialVerificationKeys() {
087                return this.verificationKeys.values();
088        }
089
090        @NonNull
091        private static McpProtectionKey copyOf(@NonNull McpProtectionKey key) {
092                byte[] keyMaterial = key.copyKeyMaterial();
093                try {
094                        return McpProtectionKey.fromIdAndBytes(key.getKeyId(), keyMaterial);
095                } finally {
096                        Arrays.fill(keyMaterial, (byte) 0);
097                }
098        }
099
100        /**
101         * Single-threaded builder for an immutable initial keyring.
102         *
103         * @author <a href="https://www.revetkn.com">Mark Allen</a>
104         */
105        @NotThreadSafe
106        public static final class Builder {
107                @NonNull
108                private final McpProtectionKey activeKey;
109                @NonNull
110                private final Map<@NonNull String, @NonNull McpProtectionKey>
111                                verificationKeys;
112
113                private Builder(@NonNull McpProtectionKey activeKey) {
114                        this.activeKey = requireNonNull(activeKey);
115                        this.verificationKeys = new LinkedHashMap<>();
116                }
117
118                /**
119                 * Adds an initial verification-only key.
120                 *
121                 * @param verificationKey key to add
122                 * @return this builder
123                 * @throws IllegalArgumentException if its ID or material duplicates an
124                 *                                  initial key
125                 */
126                @NonNull
127                public Builder addVerificationKey(
128                                @NonNull McpProtectionKey verificationKey) {
129                        addVerificationKey(this.verificationKeys,
130                                        requireNonNull(verificationKey));
131                        return this;
132                }
133
134                /**
135                 * Adds initial verification-only keys.
136                 *
137                 * @param verificationKeys keys to add
138                 * @return this builder
139                 * @throws IllegalArgumentException if an ID or material duplicates an
140                 *                                  initial key
141                 */
142                @NonNull
143                public Builder addVerificationKeys(
144                                @NonNull Collection<@NonNull McpProtectionKey> verificationKeys) {
145                        requireNonNull(verificationKeys);
146                        LinkedHashMap<@NonNull String, @NonNull McpProtectionKey> updated =
147                                        new LinkedHashMap<>(this.verificationKeys);
148                        verificationKeys.forEach(key -> addVerificationKey(updated, key));
149                        this.verificationKeys.clear();
150                        this.verificationKeys.putAll(updated);
151                        return this;
152                }
153
154                /** @return immutable initial keyring */
155                @NonNull
156                public McpProtectionKeyring build() {
157                        return new McpProtectionKeyring(this);
158                }
159
160                private void addVerificationKey(
161                                @NonNull Map<@NonNull String, @NonNull McpProtectionKey>
162                                                verificationKeys,
163                                @NonNull McpProtectionKey candidate) {
164                        requireNonNull(verificationKeys);
165                        requireNonNull(candidate);
166                        if (this.activeKey.getKeyId().equals(candidate.getKeyId())
167                                        || verificationKeys.containsKey(candidate.getKeyId()))
168                                throw new IllegalArgumentException(
169                                                "Duplicate MCP protection key ID.");
170                        if (this.activeKey.hasSameMaterial(candidate)
171                                        || verificationKeys.values().stream()
172                                        .anyMatch(existing -> existing.hasSameMaterial(candidate)))
173                                throw new IllegalArgumentException(
174                                                "Duplicate MCP protection key material.");
175                        verificationKeys.put(candidate.getKeyId(), candidate);
176                }
177        }
178}