Skip to content

Commit 0afb11f

Browse files
lu-yglu-yg
authored
feat: Add Bearer Token Support with SM4 Encryption, Improve Streaming Responses and Fix Service Error Codes (#279)
1 parent 3603e8f commit 0afb11f

File tree

7 files changed

+260
-67
lines changed

7 files changed

+260
-67
lines changed

‎base/src/main/java/com/tinyengine/it/common/exception/GlobalExceptionAdvice.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public Result<Map<String, String>> handleNullPointerException(HttpServletRequest
6969
* @param e the e
7070
* @return the result
7171
*/
72-
@ResponseStatus(HttpStatus.OK)
72+
@ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
7373
@ExceptionHandler(ServiceException.class)
7474
public Result<Map<String, String>> handleServiceException(ServiceException e) {
7575
// 修改为 log.error，传递异常对象以打印堆栈信息

‎base/src/main/java/com/tinyengine/it/common/utils/SM4Utils.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package com.tinyengine.it.common.utils;
2+
3+
import org.bouncycastle.jce.provider.BouncyCastleProvider;
4+
import javax.crypto.Cipher;
5+
import javax.crypto.KeyGenerator;
6+
import javax.crypto.SecretKey;
7+
import javax.crypto.spec.SecretKeySpec;
8+
import java.security.SecureRandom;
9+
import java.security.Security;
10+
import java.util.Base64;
11+
12+
public class SM4Utils {
13+
14+
static {
15+
Security.addProvider(new BouncyCastleProvider());
16+
}
17+
18+
private static final String ALGORITHM = "SM4";
19+
private static final String TRANSFORMATION_ECB = "SM4/ECB/PKCS5Padding";
20+
private static final int KEY_SIZE = 128;
21+
22+
/**
23+
* 生成 SM4 密钥
24+
*/
25+
public static String generateKeyBase64() throws Exception {
26+
byte[] key = generateKey();
27+
return Base64.getEncoder().encodeToString(key);
28+
}
29+
30+
public static byte[] generateKey() throws Exception {
31+
KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM, "BC");
32+
kg.init(KEY_SIZE, new SecureRandom());
33+
SecretKey secretKey = kg.generateKey();
34+
return secretKey.getEncoded();
35+
}
36+
37+
/**
38+
* ECB 模式加密 - 只加密API密钥值 (Base64 结果)
39+
*/
40+
public static String encryptECB(String apiKey, String base64Key) throws Exception {
41+
byte[] key = Base64.getDecoder().decode(base64Key);
42+
byte[] encrypted = encryptECB(apiKey.getBytes("UTF-8"), key);
43+
return Base64.getEncoder().encodeToString(encrypted);
44+
}
45+
46+
/**
47+
* ECB 模式解密 - 直接返回API密钥
48+
*/
49+
public static String decryptECB(String encryptedBase64, String base64Key) throws Exception {
50+
byte[] key = Base64.getDecoder().decode(base64Key);
51+
byte[] encrypted = Base64.getDecoder().decode(encryptedBase64);
52+
byte[] decrypted = decryptECB(encrypted, key);
53+
return new String(decrypted, "UTF-8");
54+
}
55+
56+
// ECB 模式的底层方法保持不变
57+
private static byte[] encryptECB(byte[] data, byte[] key) throws Exception {
58+
SecretKeySpec secretKeySpec = new SecretKeySpec(key, ALGORITHM);
59+
Cipher cipher = Cipher.getInstance(TRANSFORMATION_ECB, "BC");
60+
cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
61+
return cipher.doFinal(data);
62+
}
63+
64+
private static byte[] decryptECB(byte[] encryptedData, byte[] key) throws Exception {
65+
SecretKeySpec secretKeySpec = new SecretKeySpec(key, ALGORITHM);
66+
Cipher cipher = Cipher.getInstance(TRANSFORMATION_ECB, "BC");
67+
cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
68+
return cipher.doFinal(encryptedData);
69+
}
70+
71+
}

‎base/src/main/java/com/tinyengine/it/controller/AiChatController.java‎

Lines changed: 58 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,10 @@
1212

1313
package com.tinyengine.it.controller;
1414

15+
import com.tinyengine.it.common.base.Result;
16+
import com.tinyengine.it.common.exception.ExceptionEnum;
1517
import com.tinyengine.it.common.log.SystemControllerLog;
18+
import com.tinyengine.it.model.dto.AiToken;
1619
import com.tinyengine.it.model.dto.ChatRequest;
1720

1821
import com.tinyengine.it.service.app.v1.AiChatV1Service;
@@ -24,7 +27,6 @@
2427
import io.swagger.v3.oas.annotations.tags.Tag;
2528

2629
import org.springframework.beans.factory.annotation.Autowired;
27-
import org.springframework.http.HttpStatus;
2830
import org.springframework.http.MediaType;
2931
import org.springframework.http.ResponseEntity;
3032
import org.springframework.validation.annotation.Validated;
@@ -50,6 +52,7 @@ public class AiChatController {
5052
*/
5153
@Autowired
5254
private AiChatV1Service aiChatV1Service;
55+
5356
/**
5457
* AI api
5558
*
@@ -66,23 +69,29 @@ public class AiChatController {
6669
})
6770
@SystemControllerLog(description = "AI chat")
6871
@PostMapping("/ai/chat")
69-
public ResponseEntity<?> aiChat(@RequestBody ChatRequest request) {
70-
try {
71-
Object response = aiChatV1Service.chatCompletion(request);
72+
public ResponseEntity<?> aiChat(@RequestBody ChatRequest request,
73+
@RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
7274

73-
if (request.isStream()) {
74-
return ResponseEntity.ok()
75-
.contentType(MediaType.TEXT_EVENT_STREAM)
76-
.body((StreamingResponseBody) response);
77-
} else {
78-
return ResponseEntity.ok(response);
79-
}
80-
} catch (Exception e) {
81-
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
82-
.body(e.getMessage());
75+
if (authorization != null && authorization.startsWith("Bearer ")) {
76+
String token = authorization.replace("Bearer ", "");
77+
request.setApiKey(token);
8378
}
79+
80+
Object response = aiChatV1Service.chatCompletion(request);
81+
82+
if (request.isStream()) {
83+
return ResponseEntity.ok()
84+
.contentType(MediaType.TEXT_EVENT_STREAM)
85+
.header("Cache-Control", "no-cache")
86+
.header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
87+
.body((StreamingResponseBody) response);
88+
} else {
89+
return ResponseEntity.ok(response);
90+
}
91+
8492
}
8593

94+
8695
/**
8796
* AI api v1
8897
*
@@ -100,24 +109,46 @@ public ResponseEntity<?> aiChat(@RequestBody ChatRequest request) {
100109
@SystemControllerLog(description = "AI completions")
101110
@PostMapping("/chat/completions")
102111
public ResponseEntity<?> completions(@RequestBody ChatRequest request,
103-
@RequestHeader("Authorization") String authorization) {
112+
@RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
104113
if (authorization != null && authorization.startsWith("Bearer ")) {
105114
String token = authorization.replace("Bearer ", "");
106115
request.setApiKey(token);
107116
}
108-
try {
109-
Object response = aiChatV1Service.chatCompletion(request);
110117

111-
if (request.isStream()) {
112-
return ResponseEntity.ok()
113-
.contentType(MediaType.TEXT_EVENT_STREAM)
114-
.body((StreamingResponseBody) response);
115-
} else {
116-
return ResponseEntity.ok(response);
117-
}
118-
} catch (Exception e) {
119-
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
120-
.body(e.getMessage());
118+
Object response = aiChatV1Service.chatCompletion(request);
119+
120+
if (request.isStream()) {
121+
return ResponseEntity.ok()
122+
.contentType(MediaType.TEXT_EVENT_STREAM)
123+
.header("Cache-Control", "no-cache")
124+
.header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
125+
.body((StreamingResponseBody) response);
126+
} else {
127+
return ResponseEntity.ok(response);
128+
}
129+
}
130+
/**
131+
* get token
132+
*
133+
* @param request the request
134+
* @return ai回答信息 result
135+
*/
136+
@Operation(summary = "获取加密key信息", description = "获取加密key信息",
137+
parameters = {
138+
@Parameter(name = "request", description = "入参对象")
139+
}, responses = {
140+
@ApiResponse(responseCode = "200", description = "返回信息",
141+
content = @Content(mediaType = "application/json", schema = @Schema())),
142+
@ApiResponse(responseCode = "400", description = "请求失败")
143+
})
144+
@SystemControllerLog(description = "get token")
145+
@PostMapping("/encrypt-key")
146+
public Result<AiToken> getToken(@RequestBody ChatRequest request) throws Exception {
147+
String apiKey = request.getApiKey();
148+
if(apiKey == null || apiKey.isEmpty()) {
149+
return Result.failed(ExceptionEnum.CM320);
121150
}
151+
String token = aiChatV1Service.getToken(apiKey);
152+
return Result.success(new AiToken(token));
122153
}
123154
}

‎base/src/main/java/com/tinyengine/it/model/dto/AiToken.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/**
2+
* Copyright (c) 2023 - present TinyEngine Authors.
3+
* Copyright (c) 2023 - present Huawei Cloud Computing Technologies Co., Ltd.
4+
*
5+
* Use of this source code is governed by an MIT-style license.
6+
*
7+
* THE OPEN SOURCE SOFTWARE IN THIS PRODUCT IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL,
8+
* BUT WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF MERCHANTABILITY OR FITNESS FOR
9+
* A PARTICULAR PURPOSE. SEE THE APPLICABLE LICENSES FOR MORE DETAILS.
10+
*
11+
*/
12+
13+
package com.tinyengine.it.model.dto;
14+
15+
import lombok.Data;
16+
17+
/**
18+
* The type Ai token.
19+
*
20+
* @since 2025-11-27
21+
*/
22+
23+
@Data
24+
public class AiToken {
25+
String token;
26+
27+
public AiToken(String token) {
28+
this.token = token;
29+
}
30+
public AiToken(){
31+
32+
}
33+
}

0 commit comments

Comments
 (0)