Skip to content

Commit 7e218cc

Browse files
dabai827dabai827
committed
feat: (1)去掉无关日志打印 （2）去掉通过环境变量获取证书功能 （3）支持图片格式的加密
1 parent 886e3f7 commit 7e218cc

File tree

6 files changed

+206
-269
lines changed

6 files changed

+206
-269
lines changed

‎volcengine-java-sdk-ark-runtime/src/main/java/com/volcengine/ark/runtime/interceptor/EncryptionInterceptor.java‎

Lines changed: 109 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
import com.fasterxml.jackson.core.type.TypeReference;
44
import com.fasterxml.jackson.databind.ObjectMapper;
55
import com.volcengine.ark.runtime.service.CertificateManager;
6-
import com.volcengine.ark.runtime.utils.LoggerUtil;
76
import okhttp3.*;
87
import okio.Buffer;
98

9+
import java.net.URI;
10+
import java.net.URISyntaxException;
11+
1012
import org.bouncycastle.jce.provider.BouncyCastleProvider;
1113

1214
import java.security.GeneralSecurityException;
@@ -49,25 +51,21 @@ public Response intercept(Chain chain) throws IOException {
4951
Request request = chain.request();
5052

5153
String is_encrypt = request.headers().get("x-is-encrypted");
52-
LoggerUtil.debug("is-encrypt: " + is_encrypt);
5354

5455
RequestBody originalBody = request.body();
5556
if (originalBody == null) {
56-
LoggerUtil.debug("Request body is null, proceeding with original request");
5757
return chain.proceed(request);
5858
}
5959

6060
// 读取并解析请求体
6161
Map<String, Object> requestBodyJson = parseRequestBody(originalBody);
6262
String model = requestBodyJson.get("model").toString();
63-
LoggerUtil.debug("Starting encryption process for model: " + model);
6463

6564
// 非加密模式直接处理
6665
if (!"true".equals(is_encrypt)) {
6766
return proceedWithoutEncryption(chain, request, requestBodyJson);
6867
}
6968

70-
7169
// 加密模式处理
7270
return proceedWithEncryption(chain, request, requestBodyJson, model);
7371
}
@@ -78,21 +76,16 @@ private Response proceedWithEncryption(Chain chain, Request request, Map<String,
7876
if (certInfo == null) {
7977
throw new IOException("Failed to get server certificate for encryption");
8078
}
81-
LoggerUtil.debug("Successfully obtained server certificate - RingID: " + certInfo.getRingId() + ", KeyID: " + certInfo.getKeyId());
82-
8379
// 生成会话密钥和令牌
8480
SessionData sessionData;
8581
try {
8682
sessionData = KeyAgreementUtil.generateEciesKeyPair(certInfo.getPublicKey());
8783
} catch (GeneralSecurityException e) {
88-
LoggerUtil.error("generate sessionTokenData failed:", e);
8984
throw new RuntimeException(e);
9085
}
9186
byte[] e2eKey = sessionData.getCryptoKey();
9287
byte[] e2eNonce = sessionData.getCryptoNonce();
9388
String sessionToken = sessionData.getSessionToken();
94-
LoggerUtil.debug("Generated session token and crypto keys successfully");
95-
9689
// 加密请求体
9790
RequestBody encryptedBody = encryptRequestBody(requestBodyJson, e2eKey, e2eNonce);
9891

@@ -105,10 +98,8 @@ private Response proceedWithEncryption(Chain chain, Request request, Map<String,
10598
requestBuilder.addHeader("X-Session-Token", sessionToken);
10699
Request encryptedRequest = requestBuilder.build();
107100

108-
LoggerUtil.debug("Sending encrypted request to server");
109101
Response originalResponse = chain.proceed(encryptedRequest);
110102

111-
LoggerUtil.debug("Received response from server, status code: " + originalResponse.code());
112103

113104
// 处理失败响应
114105
if (!originalResponse.isSuccessful()) {
@@ -134,7 +125,6 @@ private Map<String, Object> parseRequestBody(RequestBody body) throws IOExceptio
134125
* 非加密模式处理 - 直接转发请求
135126
*/
136127
private Response proceedWithoutEncryption(Chain chain, Request request, Map<String, Object> requestBodyJson) throws IOException {
137-
LoggerUtil.debug("plaintext mode");
138128
String modifiedRequestBodyStr = mapper.writeValueAsString(requestBodyJson);
139129
RequestBody modifiedBody = RequestBody.create(
140130
MediaType.get("application/json"),
@@ -152,43 +142,120 @@ private Response proceedWithoutEncryption(Chain chain, Request request, Map<Stri
152142
*/
153143
private RequestBody encryptRequestBody(Map<String, Object> requestBodyJson, byte[] e2eKey, byte[] e2eNonce) throws IOException {
154144
try {
155-
if (requestBodyJson.containsKey("messages")) {
156-
Object messagesObj = requestBodyJson.get("messages");
157-
if (messagesObj instanceof List) {
158-
List<?> messagesList = (List<?>) messagesObj;
159-
LoggerUtil.debug("Processing " + messagesList.size() + " messages for encryption");
160-
if (!messagesList.isEmpty()) {
161-
List<Map<String, Object>> encryptedMessages = new ArrayList<>();
162-
for (Object message : messagesList) {
163-
if (message instanceof Map) {
164-
Map<?, ?> messageMap = (Map<?, ?>) message;
165-
String role = messageMap.get("role") != null ? messageMap.get("role").toString() : "user";
166-
167-
if (messageMap.get("content") != null) {
168-
String content = messageMap.get("content").toString();
169-
String encryptedContent = encryptStringWithKey(e2eKey, e2eNonce, content);
170-
Map<String, Object> encryptedMessage = new HashMap<>();
171-
encryptedMessage.put("role", role);
172-
encryptedMessage.put("content", encryptedContent);
173-
encryptedMessages.add(encryptedMessage);
174-
}
175-
}
176-
}
177-
requestBodyJson.put("messages", encryptedMessages);
178-
LoggerUtil.debug("Successfully encrypted " + encryptedMessages.size() + " messages");
145+
Object messagesObj = requestBodyJson.get("messages");
146+
if (messagesObj instanceof List) {
147+
List<?> messagesList = (List<?>) messagesObj;
148+
List<Map<String, Object>> processedMessages = new ArrayList<>();
149+
150+
for (Object message : messagesList) {
151+
if (message instanceof Map) {
152+
@SuppressWarnings("unchecked")
153+
Map<String, Object> messageMap = (Map<String, Object>) message;
154+
processedMessages.add(processMessage(messageMap, e2eKey, e2eNonce));
179155
}
180156
}
157+
requestBodyJson.put("messages", processedMessages);
181158
}
182159

183160
String modifiedRequestBodyStr = mapper.writeValueAsString(requestBodyJson);
184161
return RequestBody.create(MediaType.get("application/json"), modifiedRequestBodyStr);
185162

186163
} catch (Exception e) {
187-
LoggerUtil.debug("Failed to process request body: " + e.getMessage());
188164
throw new IOException("Failed to process request body", e);
189165
}
190166
}
191167

168+
/**
169+
* 处理单条消息
170+
*/
171+
private Map<String, Object> processMessage(Map<String, Object> message, byte[] e2eKey, byte[] e2eNonce) throws IOException {
172+
Object content = message.get("content");
173+
if (content != null) {
174+
message.put("content", processMessageContent(content, e2eKey, e2eNonce));
175+
}
176+
return message;
177+
}
178+
179+
/**
180+
* 处理消息内容
181+
*/
182+
private Object processMessageContent(Object content, byte[] e2eKey, byte[] e2eNonce) throws IOException {
183+
if (content instanceof String) {
184+
// text
185+
return encryptStringWithKey(e2eKey, e2eNonce, (String) content);
186+
}
187+
else if (content instanceof Iterable) {
188+
// multiParts
189+
List<Object> processedParts = new ArrayList<>();
190+
for (Object part : (Iterable<?>) content) {
191+
if (part instanceof Map) {
192+
@SuppressWarnings("unchecked")
193+
Map<String, Object> partMap = (Map<String, Object>) part;
194+
processedParts.add(processContentPart(partMap, e2eKey, e2eNonce));
195+
} else {
196+
throw new IOException("encryption is not supported for content type " + part.getClass().getSimpleName());
197+
}
198+
}
199+
return processedParts;
200+
}
201+
else {
202+
throw new IOException("encryption is not supported for content type " + content.getClass().getSimpleName());
203+
}
204+
}
205+
206+
/**
207+
* 处理内容部分
208+
*/
209+
private Map<String, Object> processContentPart(Map<String, Object> part, byte[] e2eKey, byte[] e2eNonce) throws IOException {
210+
String type = part.get("type").toString();
211+
212+
switch (type) {
213+
case "text":
214+
// 加密文本
215+
part.put("text", encryptStringWithKey(e2eKey, e2eNonce, part.get("text").toString()));
216+
break;
217+
218+
case "image_url":
219+
@SuppressWarnings("unchecked")
220+
Map<String, Object> imageUrl = (Map<String, Object>) part.get("image_url");
221+
processImageUrl(imageUrl, e2eKey, e2eNonce);
222+
break;
223+
224+
default:
225+
throw new IOException("encryption is not supported for content type " + type);
226+
}
227+
228+
return part;
229+
}
230+
231+
/**
232+
* 处理图片URL
233+
*/
234+
private void processImageUrl(Map<String, Object> imageUrl, byte[] e2eKey, byte[] e2eNonce) throws IOException {
235+
String url = imageUrl.get("url").toString();
236+
try {
237+
URI uri = new URI(url);
238+
String scheme = uri.getScheme();
239+
if ("data".equals(scheme)) {
240+
// 加密data URL
241+
imageUrl.put("url", encryptStringWithKey(e2eKey, e2eNonce, url));
242+
}
243+
else if ("http".equals(scheme) || "https".equals(scheme)) {
244+
System.err.println("encryption is not supported for image url, please use base64 image if you want encryption");
245+
}
246+
else {
247+
throw new IOException("encryption is not supported for image url scheme " + scheme);
248+
}
249+
250+
} catch (URISyntaxException e) {
251+
if (url.startsWith("data:")) {
252+
// 加密data URL
253+
imageUrl.put("url", encryptStringWithKey(e2eKey, e2eNonce, url));
254+
} else {
255+
throw new IOException("Invalid image URL format: " + url, e);
256+
}
257+
}
258+
}
192259

193260
/**
194261
* 添加AICC加密信息头
@@ -203,7 +270,6 @@ private void addAiccEncryptionHeader(Request.Builder requestBuilder, Certificate
203270
info.put("KeyID", certInfo.getKeyId());
204271
String infoJson = mapper.writeValueAsString(info);
205272
requestBuilder.addHeader("X-Encrypt-Info", infoJson);
206-
LoggerUtil.debug("AICC encryption enabled, added X-Encrypt-Info header");
207273
}
208274
}
209275

@@ -214,15 +280,12 @@ private Response handleErrorResponse(Response response) throws IOException {
214280
ResponseBody errorBody = response.body();
215281
if (errorBody != null) {
216282
String errorResponseStr = errorBody.string();
217-
LoggerUtil.debug("Request failed with status: " + response.code() + ", error response: " + errorResponseStr);
218283
MediaType contentType = errorBody.contentType();
219284
if (contentType == null) {
220285
contentType = MediaType.get("application/json; charset=utf-8");
221286
}
222287
ResponseBody newErrorBody = ResponseBody.create(contentType, errorResponseStr);
223288
return response.newBuilder().body(newErrorBody).build();
224-
} else {
225-
LoggerUtil.debug("Request failed with status: " + response.code() + ", no response body");
226289
}
227290
return response;
228291
}
@@ -234,7 +297,6 @@ private Response decryptResponse(byte[] key, byte[] nonce, Response response) {
234297
try {
235298
ResponseBody responseBody = response.body();
236299
if (responseBody == null) {
237-
LoggerUtil.debug("Response body is null, returning original response");
238300
return response;
239301
}
240302

@@ -251,13 +313,11 @@ private Response decryptResponse(byte[] key, byte[] nonce, Response response) {
251313
);
252314
return handleNormalResponse(key, nonce, response, responseJson);
253315
} catch (Exception e) {
254-
LoggerUtil.debug("Failed to parse response as JSON, trying as stream: " + e.getMessage());
255316
return handleStreamResponse(key, nonce, response, responseBodyStr);
256317
}
257318
}
258319

259320
} catch (Exception e) {
260-
LoggerUtil.debug("Failed to decrypt response: " + e.getMessage());
261321
return response;
262322
}
263323
}
@@ -269,11 +329,9 @@ private Response handleStreamResponse(byte[] key, byte[] nonce, Response respons
269329
try {
270330
ResponseBody originalBody = response.body();
271331
if (originalBody == null) {
272-
LoggerUtil.debug("Original body is null in handleStreamResponse");
273332
return response;
274333
}
275334
String decryptedContent = decryptStreamContent(key, nonce, originalContent);
276-
LoggerUtil.debug("Stream response decrypted successfully, decrypted length: " + decryptedContent.length());
277335

278336
MediaType contentType = originalBody.contentType();
279337
if (contentType == null) {
@@ -289,7 +347,6 @@ private Response handleStreamResponse(byte[] key, byte[] nonce, Response respons
289347
.build();
290348

291349
} catch (Exception e) {
292-
LoggerUtil.debug("Failed to handle stream response: " + e.getMessage());
293350
return response;
294351
}
295352
}
@@ -302,14 +359,11 @@ private Response handleNormalResponse(byte[] key, byte[] nonce, Response respons
302359
@SuppressWarnings("unchecked")
303360
List<Map<String, Object>> choices = (List<Map<String, Object>>) responseJson.get("choices");
304361

305-
for (int i = 0; i < choices.size(); i++) {
306-
Map<String, Object> choice = choices.get(i);
362+
for (Map<String, Object> choice : choices) {
307363
if (shouldDecryptChoice(choice)) {
308-
decryptChoiceContent(key, nonce, choice, i);
364+
decryptChoiceContent(key, nonce, choice);
309365
}
310366
}
311-
} else {
312-
LoggerUtil.debug("No choices found in response or choices is not a list");
313367
}
314368

315369
String modifiedResponseBodyStr = mapper.writeValueAsString(responseJson);
@@ -334,25 +388,19 @@ private Response handleNormalResponse(byte[] key, byte[] nonce, Response respons
334388
/**
335389
* 解密单个choice内容
336390
*/
337-
private void decryptChoiceContent(byte[] key, byte[] nonce, Map<String, Object> choice, int index) {
391+
private void decryptChoiceContent(byte[] key, byte[] nonce, Map<String, Object> choice) {
338392
String encryptedContent = getEncryptedContentFromChoice(choice);
339-
// LoggerUtil.debug("Choice " + index + " requires decryption, encrypted content length: " +
340-
// (encryptedContent != null ? encryptedContent.length() : 0));
341-
342393
if (encryptedContent != null && !encryptedContent.isEmpty()) {
343394
try {
344395
String decryptedContent = decryptStringWithKey(key, nonce, encryptedContent);
345396
@SuppressWarnings("unchecked")
346397
Map<String, Object> message = (Map<String, Object>) choice.get("message");
347398
message.put("content", decryptedContent);
348399
} catch (Exception e) {
349-
LoggerUtil.debug("Failed to decrypt content for choice " + index + ": " + e.getMessage());
350400
@SuppressWarnings("unchecked")
351401
Map<String, Object> message = (Map<String, Object>) choice.get("message");
352402
message.put("content", "");
353403
}
354-
} else {
355-
LoggerUtil.debug("Encrypted content is null or empty for choice " + index);
356404
}
357405
}
358406

@@ -362,7 +410,6 @@ private void decryptChoiceContent(byte[] key, byte[] nonce, Map<String, Object>
362410
private String decryptStreamContent(byte[] key, byte[] nonce, String streamContent) {
363411
try {
364412
String[] lines = streamContent.split("\n");
365-
LoggerUtil.debug("Stream content split into " + lines.length + " lines");
366413
StringBuilder decryptedContent = new StringBuilder();
367414

368415
for (String line : lines) {
@@ -385,7 +432,6 @@ private String decryptStreamContent(byte[] key, byte[] nonce, String streamConte
385432
String decryptedJson = mapper.writeValueAsString(decryptedChunk);
386433
decryptedContent.append("data: ").append(decryptedJson).append("\n");
387434
} catch (Exception e) {
388-
LoggerUtil.debug("Failed to process stream data line: " + e.getMessage() + ", line: " + line);
389435
decryptedContent.append(line).append("\n");
390436
}
391437
} else {
@@ -395,7 +441,6 @@ private String decryptStreamContent(byte[] key, byte[] nonce, String streamConte
395441

396442
return decryptedContent.toString();
397443
} catch (Exception e) {
398-
LoggerUtil.debug("Failed to decrypt stream content: " + e.getMessage());
399444
return streamContent;
400445
}
401446
}
@@ -409,8 +454,7 @@ private Map<String, Object> decryptStreamChunk(byte[] key, byte[] nonce, Map<Str
409454
@SuppressWarnings("unchecked")
410455
List<Map<String, Object>> choices = (List<Map<String, Object>>) chunkData.get("choices");
411456

412-
for (int i = 0; i < choices.size(); i++) {
413-
Map<String, Object> choice = choices.get(i);
457+
for (Map<String, Object> choice : choices) {
414458
if (shouldDecryptStreamChoice(choice)) {
415459
String encryptedContent = getEncryptedContentFromStreamChoice(choice);
416460
if (encryptedContent != null && !encryptedContent.isEmpty()) {
@@ -420,7 +464,6 @@ private Map<String, Object> decryptStreamChunk(byte[] key, byte[] nonce, Map<Str
420464
Map<String, Object> delta = (Map<String, Object>) choice.get("delta");
421465
delta.put("content", decryptedContent);
422466
} catch (Exception e) {
423-
LoggerUtil.debug("Failed to decrypt stream content for choice " + i + ": " + e.getMessage());
424467
@SuppressWarnings("unchecked")
425468
Map<String, Object> delta = (Map<String, Object>) choice.get("delta");
426469
delta.put("content", "");
@@ -431,7 +474,6 @@ private Map<String, Object> decryptStreamChunk(byte[] key, byte[] nonce, Map<Str
431474
}
432475
return chunkData;
433476
} catch (Exception e) {
434-
LoggerUtil.debug("Failed to decrypt stream chunk: " + e.getMessage());
435477
return chunkData;
436478
}
437479
}

0 commit comments

Comments
 (0)