feat: Add Bearer Token Support with SM4 Encryption, Improve Streaming Responses and Fix Service Error Codes #279
Conversation
…into feat/rag # Conflicts: # base/src/main/java/com/tinyengine/it/common/utils/ImageThumbnailGenerator.java # base/src/main/java/com/tinyengine/it/controller/VectorStorageController.java # base/src/main/java/com/tinyengine/it/rag/entity/EmbeddingMatchDto.java # base/src/main/java/com/tinyengine/it/rag/entity/VectorDocument.java # base/src/main/java/com/tinyengine/it/rag/service/StorageService.java # base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
WalkthroughThe changes add SM4 key utilities and a BouncyCastle dependency, introduce token generation/consumption using an EKEY_ prefix, accept Bearer tokens in controller endpoints and add /encrypt-key, adjust streaming headers, and change ServiceException handler response code from 200 to 500. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Controller as AiChatController
participant Service as AiChatV1Service
participant SM4 as SM4Utils
participant Response
rect rgb(220,240,220)
Note over Client,Service: Token generation (/encrypt-key)
Client->>Controller: POST /encrypt-key { apiKey }
Controller->>Service: getToken(apiKey)
Service->>SM4: encryptECB(apiKey)
SM4-->>Service: encrypted(Base64)
Service-->>Controller: "EKEY_<encrypted>"
Controller-->>Response: Result<AiToken>
end
rect rgb(240,240,220)
Note over Client,Response: Bearer usage in aiChat/completions
Client->>Controller: POST /aiChat (Authorization: Bearer EKEY_<token>)
Controller->>Service: aiChat(request, bearerToken)
Service->>SM4: if startsWith("EKEY_") decryptECB(<encrypted part>)
SM4-->>Service: decrypted apiKey
Service->>Service: processRequest(apiKey)
alt Stream Response
Controller->>Response: Cache-Control:no-cache, X-Accel-Buffering:no, stream data
else Standard Response
Controller->>Response: standard response
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
pom.xml (1)
37-42: Remove duplicate property declaration.
bcprov-jdk18on.versionis declared twice (lines 37 and 41). Remove one to avoid confusion.🔎 Proposed fix
<bcprov-jdk18on.version>1.79</bcprov-jdk18on.version> <jjwt.version>0.12.3</jjwt.version> <langchain4j.version>1.5.0</langchain4j.version> <langchain4j-rag.version>1.5.0-beta11</langchain4j-rag.version> - <bcprov-jdk18on.version>1.79</bcprov-jdk18on.version> <tiny-engine.version>1.0-SNAPSHOT</tiny-engine.version>base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java (2)
137-140:temperatureis set twice unconditionally then conditionally.Line 137 sets
temperatureunconditionally (even if null), then line 138-140 sets it again conditionally. Remove the unconditional assignment.🔎 Proposed fix
- body.put("temperature", request.getTemperature()); if (request.getTemperature() != null) { body.put("temperature", request.getTemperature()); }
153-155: Incorrect null check—usesgetMaxInputTokens()but putsgetVlHighResolutionImages().This appears to be a copy/paste error. The condition should check
getVlHighResolutionImages().🔎 Proposed fix
- if (request.getMaxInputTokens() != null) { + if (request.getVlHighResolutionImages() != null) { body.put("vl_high_resolution_images", request.getVlHighResolutionImages()); }
🧹 Nitpick comments (6)
base/src/main/java/com/tinyengine/it/model/dto/AiToken.java (1)
23-33: Field should beprivate; consider using Lombok constructor annotations.The
tokenfield lacks an access modifier (defaulting to package-private). For proper encapsulation, declare itprivate. Also, the explicit constructors can be replaced with Lombok annotations.🔎 Proposed refactor
@Data +@NoArgsConstructor +@AllArgsConstructor public class AiToken { - String token; - - public AiToken(String token) { - this.token = token; - } - public AiToken(){ - - } + private String token; }Add the imports:
import lombok.AllArgsConstructor; import lombok.NoArgsConstructor;base/src/main/java/com/tinyengine/it/common/utils/SM4Utils.java (3)
1-11: Missing copyright header.This file lacks the standard copyright header present in other files in the project.
🔎 Add copyright header
+/** + * Copyright (c) 2023 - present TinyEngine Authors. + * Copyright (c) 2023 - present Huawei Cloud Computing Technologies Co., Ltd. + * + * Use of this source code is governed by an MIT-style license. + * + * THE OPEN SOURCE SOFTWARE IN THIS PRODUCT IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL, + * BUT WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF MERCHANTABILITY OR FITNESS FOR + * A PARTICULAR PURPOSE. SEE THE APPLICABLE LICENSES FOR MORE DETAILS. + * + */ + package com.tinyengine.it.common.utils;
18-20: ECB mode lacks semantic security; consider CBC or GCM for stronger protection.The static analysis warnings about DES are false positives (SM4 is China's national block cipher, not DES). However, the ECB mode concern is valid—ECB is deterministic and can leak patterns for repeated plaintext blocks.
For short, unique API keys, ECB may be acceptable, but CBC or GCM modes provide better security guarantees. If you decide to upgrade, use
SM4/CBC/PKCS5Paddingwith a random IV orSM4/GCM/NoPadding.
40-54: UseStandardCharsets.UTF_8instead of string literal.Replace
"UTF-8"withStandardCharsets.UTF_8for type safety and to avoid potentialUnsupportedEncodingException.🔎 Proposed fix
public static String encryptECB(String apiKey, String base64Key) throws Exception { byte[] key = Base64.getDecoder().decode(base64Key); - byte[] encrypted = encryptECB(apiKey.getBytes("UTF-8"), key); + byte[] encrypted = encryptECB(apiKey.getBytes(StandardCharsets.UTF_8), key); return Base64.getEncoder().encodeToString(encrypted); } // ... public static String decryptECB(String encryptedBase64, String base64Key) throws Exception { byte[] key = Base64.getDecoder().decode(base64Key); byte[] encrypted = Base64.getDecoder().decode(encryptedBase64); byte[] decrypted = decryptECB(encrypted, key); - return new String(decrypted, "UTF-8"); + return new String(decrypted, StandardCharsets.UTF_8); }base/src/main/java/com/tinyengine/it/controller/AiChatController.java (2)
112-130: Duplicate code withaiChatmethod.Lines 114-129 duplicate the Bearer token extraction and streaming response logic from
aiChat. Consider extracting common logic into private helper methods to improve maintainability.🔎 Example refactor
private void extractBearerToken(String authorization, ChatRequest request) { if (authorization != null && authorization.startsWith("Bearer ")) { request.setApiKey(authorization.substring(7)); } } private ResponseEntity<?> buildChatResponse(Object response, boolean isStream) { if (isStream) { return ResponseEntity.ok() .contentType(MediaType.TEXT_EVENT_STREAM) .header("Cache-Control", "no-cache") .header("X-Accel-Buffering", "no") .body((StreamingResponseBody) response); } return ResponseEntity.ok(response); }
149-149: Minor formatting: missing space afterif.🔎 Proposed fix
- if(apiKey == null || apiKey.isEmpty()) { + if (apiKey == null || apiKey.isEmpty()) {
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
base/src/main/java/com/tinyengine/it/common/exception/GlobalExceptionAdvice.java(1 hunks)base/src/main/java/com/tinyengine/it/common/utils/SM4Utils.java(1 hunks)base/src/main/java/com/tinyengine/it/controller/AiChatController.java(4 hunks)base/src/main/java/com/tinyengine/it/model/dto/AiToken.java(1 hunks)base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java(6 hunks)base/src/main/java/com/tinyengine/it/service/app/v1/AiChatV1Service.java(1 hunks)pom.xml(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java (3)
base/src/main/java/com/tinyengine/it/common/utils/JsonUtils.java (1)
JsonUtils(45-349)base/src/main/java/com/tinyengine/it/common/utils/SM4Utils.java (1)
SM4Utils(12-71)base/src/main/java/com/tinyengine/it/common/utils/Utils.java (1)
Slf4j(59-517)
🪛 ast-grep (0.40.0)
base/src/main/java/com/tinyengine/it/common/utils/SM4Utils.java
[warning] 30-30: Triple DES (3DES or DESede) is considered deprecated. AES is the recommended cipher. Upgrade to use AES.
Context: KeyGenerator.getInstance(ALGORITHM, "BC")
Note: [CWE-326]: Inadequate Encryption Strength [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://find-sec-bugs.github.io/bugs.htm#TDES_USAGE
- https://csrc.nist.gov/News/2017/Update-to-Current-Use-and-Deprecation-of-TDEA
(desede-is-deprecated-java)
[warning] 58-58: Triple DES (3DES or DESede) is considered deprecated. AES is the recommended cipher. Upgrade to use AES.
Context: Cipher.getInstance(TRANSFORMATION_ECB, "BC")
Note: [CWE-326]: Inadequate Encryption Strength [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://find-sec-bugs.github.io/bugs.htm#TDES_USAGE
- https://csrc.nist.gov/News/2017/Update-to-Current-Use-and-Deprecation-of-TDEA
(desede-is-deprecated-java)
[warning] 65-65: Triple DES (3DES or DESede) is considered deprecated. AES is the recommended cipher. Upgrade to use AES.
Context: Cipher.getInstance(TRANSFORMATION_ECB, "BC")
Note: [CWE-326]: Inadequate Encryption Strength [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://find-sec-bugs.github.io/bugs.htm#TDES_USAGE
- https://csrc.nist.gov/News/2017/Update-to-Current-Use-and-Deprecation-of-TDEA
(desede-is-deprecated-java)
[warning] 30-30: Use of AES with ECB mode detected. ECB doesn't provide message confidentiality and is not semantically secure so should not be used. Instead, use a strong, secure cipher: Cipher.getInstance("AES/CBC/PKCS7PADDING"). See https://owasp.org/www-community/Using_the_Java_Cryptographic_Extensions for more information.
Context: KeyGenerator.getInstance(ALGORITHM, "BC")
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://owasp.org/Top10/A02_2021-Cryptographic_Failures
- https://googleprojectzero.blogspot.com/2022/10/rc4-is-still-considered-harmful.html
(use-of-aes-ecb-java)
[warning] 58-58: Use of AES with ECB mode detected. ECB doesn't provide message confidentiality and is not semantically secure so should not be used. Instead, use a strong, secure cipher: Cipher.getInstance("AES/CBC/PKCS7PADDING"). See https://owasp.org/www-community/Using_the_Java_Cryptographic_Extensions for more information.
Context: Cipher.getInstance(TRANSFORMATION_ECB, "BC")
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://owasp.org/Top10/A02_2021-Cryptographic_Failures
- https://googleprojectzero.blogspot.com/2022/10/rc4-is-still-considered-harmful.html
(use-of-aes-ecb-java)
[warning] 65-65: Use of AES with ECB mode detected. ECB doesn't provide message confidentiality and is not semantically secure so should not be used. Instead, use a strong, secure cipher: Cipher.getInstance("AES/CBC/PKCS7PADDING"). See https://owasp.org/www-community/Using_the_Java_Cryptographic_Extensions for more information.
Context: Cipher.getInstance(TRANSFORMATION_ECB, "BC")
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
- https://owasp.org/Top10/A02_2021-Cryptographic_Failures
- https://googleprojectzero.blogspot.com/2022/10/rc4-is-still-considered-harmful.html
(use-of-aes-ecb-java)
🔇 Additional comments (5)
base/src/main/java/com/tinyengine/it/common/exception/GlobalExceptionAdvice.java (1)
72-78: Appropriate status code correction for ServiceException.Returning 500 for service exceptions is correct. However, the
NullPointerExceptionhandler (lines 58-64) still returns 200 OK, which is inconsistent—NPEs are typically server errors and should also return 500.Consider aligning the
NullPointerExceptionhandler to also returnHttpStatus.INTERNAL_SERVER_ERRORfor consistency.base/src/main/java/com/tinyengine/it/service/app/v1/AiChatV1Service.java (1)
31-37: New interface method looks good.The
getTokenmethod signature is consistent with the existingchatCompletionmethod. Consider using a more specific exception type in the future for better error handling, but this is acceptable for now given the existing pattern.base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java (1)
110-114: Unclear semantics for URLs ending with#.The special handling for URLs ending with
#is not documented and the purpose is unclear. Consider adding a comment explaining this behavior or reconsider if this is the intended design.base/src/main/java/com/tinyengine/it/controller/AiChatController.java (2)
146-153: Consider rate limiting for the encryption endpoint.The
/encrypt-keyendpoint generates encrypted tokens from API keys. Without rate limiting, this could be abused for:
- Enumeration attacks if error responses differ based on key validity
- Resource exhaustion through repeated encryption operations
Consider adding rate limiting or authentication requirements to this endpoint.
Is this endpoint intended to be publicly accessible without authentication? If users can call this endpoint freely, verify that rate limiting is applied at the gateway/infrastructure level.
83-91: Good streaming response headers.The
Cache-Control: no-cacheandX-Accel-Buffering: noheaders are appropriately set for Server-Sent Events, ensuring proper streaming behavior with Nginx and preventing unwanted caching.
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
Show resolved
Hide resolved
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
Show resolved
Hide resolved
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
Show resolved
Hide resolved
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
Show resolved
Hide resolved
base/src/main/java/com/tinyengine/it/service/app/impl/v1/AiChatV1ServiceImpl.java
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (4)
base/src/main/java/com/tinyengine/it/controller/AiChatController.java (4)
75-78: Usesubstringinstead ofreplacefor Bearer token extraction.This issue was already flagged in a previous review. The
replace("Bearer ", "")approach replaces all occurrences of "Bearer " in the string, not just the prefix. Useauthorization.substring(7)instead.
82-90: Unsafe cast toStreamingResponseBodywithout type verification.This issue was already flagged in a previous review. The response from
chatCompletionis cast toStreamingResponseBodywithout runtime type verification, which can causeClassCastExceptionif the service doesn't honor the contract.
113-116: Usesubstringinstead ofreplacefor Bearer token extraction.Same issue as in
aiChatmethod. This was already flagged in a previous review for theaiChatmethod.
120-128: Unsafe cast toStreamingResponseBodywithout type verification.Same issue as in
aiChatmethod. This was already flagged in a previous review for theaiChatmethod.
🧹 Nitpick comments (4)
base/src/main/java/com/tinyengine/it/controller/AiChatController.java (4)
72-73: Avoid declaring genericExceptionin throws clause.The method signature declares
throws Exception, which is too broad. Consider throwing specific exceptions (e.g.,ServiceException,IllegalArgumentException) or using Spring's exception handling mechanisms to provide clearer error semantics.
111-112: Avoid declaring genericExceptionin throws clause.Same issue as in
aiChat: the method declaresthrows Exception, which is too broad. Use specific exception types or Spring's exception handling.
146-146: Avoid declaring genericExceptionin throws clause.Same issue as the other methods: declaring
throws Exceptionis too broad. Use specific exception types.
147-149: Consider usingisBlank()or trimming the input.The validation checks for
nullorisEmpty(), but doesn't handle whitespace-only strings. Consider usingStringUtils.isBlank(apiKey)(if Apache Commons is available) orapiKey.trim().isEmpty()for more robust validation.🔎 Proposed enhancement
- if(apiKey == null || apiKey.isEmpty()) { + if(apiKey == null || apiKey.trim().isEmpty()) { return Result.failed(ExceptionEnum.CM320); }
hexqi
changed the title
2 participants
English | 简体中文
PR
PR Checklist
Please check if your PR fulfills the following requirements:
PR Type
What kind of change does this PR introduce?
Background and solution
What is the current behavior?
Issue Number: N/A
What is the new behavior?
Does this PR introduce a breaking change?
Other information
Summary by CodeRabbit
Bug Fixes
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.