1111import com .dylibso .chicory .wasm .types .FunctionType ;
1212import com .dylibso .chicory .wasm .types .ValType ;
1313import com .google .protobuf .ByteString ;
14- import com .google .protobuf .BytesValue ;
15- import com .google .protobuf .GeneratedMessageV3 ;
14+ import com .google .protobuf .GeneratedMessage ;
1615import com .google .protobuf .InvalidProtocolBufferException ;
1716import com .google .protobuf .Timestamp ;
1817import com .spotify .confidence .flags .resolver .v1 .LogMessage ;
2524import com .spotify .confidence .wasm .Messages ;
2625import java .io .IOException ;
2726import java .io .InputStream ;
28- import java .security .SecureRandom ;
2927import java .time .Instant ;
3028import java .util .List ;
3129import java .util .function .Function ;
32- import javax .crypto .Cipher ;
33- import javax .crypto .spec .IvParameterSpec ;
34- import javax .crypto .spec .SecretKeySpec ;
35- import rust_guest .Types ;
3630
3731@ FunctionalInterface
3832interface WasmFlagLogger {
@@ -82,11 +76,6 @@ public WasmResolveApi(WasmFlagLogger flagLogger, RetryStrategy retryStrategy) {
8276 "wasm_msg_current_thread_id" ,
8377 FunctionType .of (List .of (), List .of (ValType .I32 )),
8478 this ::currentThreadId ))
85- .addFunction (
86- createImportFunction (
87- "encrypt_resolve_token" ,
88- Types .EncryptionRequest ::parseFrom ,
89- this ::encryptResolveToken ))
9079 .build ())
9180 .withMachineFactory (MachineFactoryCompiler ::compile )
9281 .build ();
@@ -101,38 +90,11 @@ public WasmResolveApi(WasmFlagLogger flagLogger, RetryStrategy retryStrategy) {
10190 }
10291 }
10392
104- private GeneratedMessageV3 log (LogMessage message ) {
93+ private GeneratedMessage log (LogMessage message ) {
10594 System .out .println (message .getMessage ());
10695 return Messages .Void .getDefaultInstance ();
10796 }
10897
109- private GeneratedMessageV3 encryptResolveToken (Types .EncryptionRequest encryptionRequest ) {
110- try {
111- final byte [] tokenData = encryptionRequest .getTokenData ().toByteArray ();
112- final byte [] encryptionKey = encryptionRequest .getEncryptionKey ().toByteArray ();
113- if (encryptionKey .length != 16 ) {
114- throw new IllegalArgumentException ("Encryption key must be exactly 16 bytes for AES-128" );
115- }
116- final byte [] iv = new byte [16 ];
117- new SecureRandom ().nextBytes (iv );
118- final Cipher cipher = Cipher .getInstance ("AES/CBC/PKCS5Padding" );
119- final SecretKeySpec secretKey = new SecretKeySpec (encryptionKey , "AES" );
120- final IvParameterSpec ivSpec = new IvParameterSpec (iv );
121-
122- cipher .init (Cipher .ENCRYPT_MODE , secretKey , ivSpec );
123- final byte [] encryptedData = cipher .doFinal (tokenData );
124-
125- final byte [] result = new byte [iv .length + encryptedData .length ];
126- System .arraycopy (iv , 0 , result , 0 , iv .length );
127- System .arraycopy (encryptedData , 0 , result , iv .length , encryptedData .length );
128-
129- return BytesValue .newBuilder ().setValue (ByteString .copyFrom (result )).build ();
130-
131- } catch (Exception e ) {
132- throw new RuntimeException ("Failed to encrypt resolve token" , e );
133- }
134- }
135-
13698 private long [] currentThreadId (Instance instance , long ... longs ) {
13799 return new long [] {0 };
138100 }
@@ -185,7 +147,7 @@ public ResolveFlagsResponse resolve(ResolveFlagsRequest request) {
185147 "resolve" );
186148 }
187149
188- private <T extends GeneratedMessageV3 > T consumeResponse (int addr , ParserFn <T > codec ) {
150+ private <T extends GeneratedMessage > T consumeResponse (int addr , ParserFn <T > codec ) {
189151 try {
190152 final Messages .Response response = Messages .Response .parseFrom (consume (addr ));
191153 if (response .hasError ()) {
@@ -198,7 +160,7 @@ private <T extends GeneratedMessageV3> T consumeResponse(int addr, ParserFn<T> c
198160 }
199161 }
200162
201- private <T extends GeneratedMessageV3 > T consumeRequest (int addr , ParserFn <T > codec ) {
163+ private <T extends GeneratedMessage > T consumeRequest (int addr , ParserFn <T > codec ) {
202164 try {
203165 final Messages .Request request = Messages .Request .parseFrom (consume (addr ));
204166 return codec .apply (request .getData ().toByteArray ());
@@ -207,13 +169,13 @@ private <T extends GeneratedMessageV3> T consumeRequest(int addr, ParserFn<T> co
207169 }
208170 }
209171
210- private int transferRequest (GeneratedMessageV3 message ) {
172+ private int transferRequest (GeneratedMessage message ) {
211173 final byte [] request =
212174 Messages .Request .newBuilder ().setData (message .toByteString ()).build ().toByteArray ();
213175 return transfer (request );
214176 }
215177
216- private int transferResponseSuccess (GeneratedMessageV3 response ) {
178+ private int transferResponseSuccess (GeneratedMessage response ) {
217179 final byte [] wrapperBytes =
218180 Messages .Response .newBuilder ().setData (response .toByteString ()).build ().toByteArray ();
219181 return transfer (wrapperBytes );
@@ -240,16 +202,16 @@ private int transfer(byte[] data) {
240202 return addr ;
241203 }
242204
243- private <T extends GeneratedMessageV3 > ImportFunction createImportFunction (
244- String name , ParserFn <T > reqCodec , Function <T , GeneratedMessageV3 > impl ) {
205+ private <T extends GeneratedMessage > ImportFunction createImportFunction (
206+ String name , ParserFn <T > reqCodec , Function <T , GeneratedMessage > impl ) {
245207 return new ImportFunction (
246208 "wasm_msg" ,
247209 "wasm_msg_host_" + name ,
248210 HOST_FN_TYPE ,
249211 (instance1 , args ) -> {
250212 try {
251213 final T message = consumeRequest ((int ) args [0 ], reqCodec );
252- final GeneratedMessageV3 response = impl .apply (message );
214+ final GeneratedMessage response = impl .apply (message );
253215 return new long [] {transferResponseSuccess (response )};
254216 } catch (Exception e ) {
255217 return new long [] {transferResponseError (e .getMessage ())};
0 commit comments