Skip to content

Commit 79eb4c3

Browse files
luminghao-bytedanceluminghao-bytedance
committed
feat(ark_runtime_bot_api): add bot chat api client methods
1 parent c7a066e commit 79eb4c3

27 files changed

+858
-171
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@ public class Const {
77
public static final String CLIENT_REQUEST_HEADER = "X-Client-Request-Id";
88
public static final String SERVER_REQUEST_HEADER = "X-Request-Id";
99
public static final String REQUEST_MODEL = "X-Request-Model";
10+
public static final String REQUEST_BOT = "X-Request-Bot";
11+
public static final String REQUEST_BOT_ID = "botId";
1012
public static final Integer DEFAULT_MANDATORY_REFRESH_TIMEOUT = 10 * 60; // 10 min
1113
public static final Integer DEFAULT_ADVISORY_REFRESH_TIMEOUT = 30 * 60; // 30 min
1214
public static final Integer DEFAULT_STS_TIMEOUT = 7 * 24 * 60 * 60; // 7 days
15+
16+
public static final String RESOURCE_TYPE_BOT = "bot";
17+
public static final String RESOURCE_TYPE_ENDPOINT = "endpoint";
1318
}

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

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
package com.volcengine.ark.runtime.interceptor;
2+
3+
import com.volcengine.ApiClient;
4+
import com.volcengine.ApiException;
5+
import com.volcengine.StringUtil;
6+
import com.volcengine.ark.ArkApi;
7+
import com.volcengine.ark.model.GetApiKeyRequest;
8+
import com.volcengine.ark.model.GetApiKeyResponse;
9+
import com.volcengine.ark.runtime.Const;
10+
import com.volcengine.ark.runtime.exception.ArkException;
11+
import com.volcengine.sign.Credentials;
12+
import okhttp3.Interceptor;
13+
import okhttp3.Request;
14+
import okhttp3.Response;
15+
import org.apache.commons.lang.StringUtils;
16+
17+
import java.io.IOException;
18+
import java.util.ArrayList;
19+
import java.util.List;
20+
import java.util.Map;
21+
import java.util.Objects;
22+
import java.util.concurrent.ConcurrentHashMap;
23+
import java.util.concurrent.locks.ReentrantReadWriteLock;
24+
import java.util.function.BiFunction;
25+
26+
public class ArkResourceStsAuthenticationInterceptor implements Interceptor {
27+
28+
private final String ak;
29+
private final String sk;
30+
// cacheKey = resourceType/resourceId
31+
private Map<String, ArkResourceStsTokenInfo> resourceStsTokens;
32+
private final Integer advisoryRefreshTimeout = Const.DEFAULT_ADVISORY_REFRESH_TIMEOUT;
33+
private final Integer mandatoryRefreshTimeout = Const.DEFAULT_MANDATORY_REFRESH_TIMEOUT;
34+
private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
35+
private final ArkApi volcClient;
36+
37+
public ArkResourceStsAuthenticationInterceptor(String ak, String sk, String region) {
38+
Objects.requireNonNull(ak, "Ak token required");
39+
Objects.requireNonNull(sk, "Sk token required");
40+
this.ak = ak;
41+
this.sk = sk;
42+
this.resourceStsTokens = new ConcurrentHashMap<>();
43+
44+
ApiClient apiClient = new ApiClient()
45+
.setCredentials(Credentials.getCredentials(ak, sk))
46+
.setRegion(region);
47+
ArkApi arkApi = new ArkApi(apiClient);
48+
this.volcClient = arkApi;
49+
}
50+
51+
@Override
52+
public Response intercept(Chain chain) throws IOException {
53+
Request request = chain.request();
54+
String requestResourceType = getRequestResourceType(request);
55+
String requestResourceId = getRequestResourceId(request);
56+
57+
Request newRequest = chain.request()
58+
.newBuilder()
59+
.header("Authorization", "Bearer " + getResourceStsToken(requestResourceType, requestResourceId))
60+
.build();
61+
return chain.proceed(newRequest);
62+
}
63+
64+
private String getRequestResourceType(Request request) {
65+
if (StringUtils.isNotBlank(request.header(Const.REQUEST_BOT))) {
66+
return Const.RESOURCE_TYPE_BOT;
67+
}
68+
return Const.RESOURCE_TYPE_ENDPOINT;
69+
}
70+
71+
private String getRequestResourceId(Request request) {
72+
if (StringUtils.isNotBlank(request.header(Const.REQUEST_BOT))) {
73+
return request.header(Const.REQUEST_BOT);
74+
}
75+
return request.header(Const.REQUEST_MODEL);
76+
}
77+
78+
private String getResourceStsToken(String resourceType, String resourceId) {
79+
refresh(resourceType, resourceId);
80+
81+
ArkResourceStsTokenInfo tokenInfo = this.resourceStsTokens.get(getResourceKey(resourceType, resourceId));
82+
if (tokenInfo == null) {
83+
return "";
84+
}
85+
return tokenInfo.getToken();
86+
}
87+
88+
private void refresh(String resourceType, String resourceId) {
89+
if (!need_refresh(resourceType, resourceId, this.advisoryRefreshTimeout)) {
90+
return;
91+
}
92+
93+
if (lock.writeLock().tryLock()) {
94+
if (!need_refresh(resourceType, resourceId, this.advisoryRefreshTimeout)) {
95+
return;
96+
}
97+
98+
try {
99+
boolean isMandatoryRefresh = need_refresh(resourceType, resourceId, this.mandatoryRefreshTimeout);
100+
protectedRefresh(resourceType, resourceId, isMandatoryRefresh);
101+
} finally {
102+
lock.writeLock().unlock();
103+
}
104+
} else if (need_refresh(resourceType, resourceId, this.mandatoryRefreshTimeout)) {
105+
try {
106+
lock.writeLock().lock();
107+
if (!need_refresh(resourceType, resourceId, this.mandatoryRefreshTimeout)) {
108+
return;
109+
}
110+
protectedRefresh(resourceType, resourceId, true);
111+
} finally {
112+
lock.writeLock().unlock();
113+
}
114+
}
115+
}
116+
117+
private boolean need_refresh(String resourceType, String resourceId, Integer refresh_in) {
118+
ArkResourceStsTokenInfo tokenInfo = this.resourceStsTokens.get(getResourceKey(resourceType, resourceId));
119+
if (tokenInfo == null) {
120+
return true;
121+
}
122+
123+
return tokenInfo.getExpiredTime() - System.currentTimeMillis() / 1000 < refresh_in;
124+
}
125+
126+
private void protectedRefresh(String resourceType, String resourceId, boolean isMandatory) {
127+
this.resourceStsTokens.compute(getResourceKey(resourceType, resourceId), new BiFunction<String, ArkResourceStsTokenInfo, ArkResourceStsTokenInfo>() {
128+
@Override
129+
public ArkResourceStsTokenInfo apply(String s, ArkResourceStsTokenInfo stringIntegerPair) {
130+
try {
131+
ArkResourceStsTokenInfo tokenInfo = getToken(resourceType, resourceId, Const.DEFAULT_STS_TIMEOUT);
132+
return tokenInfo;
133+
} catch (ApiException e) {
134+
if (isMandatory) {
135+
throw new RuntimeException(e);
136+
}
137+
return null;
138+
}
139+
}
140+
});
141+
}
142+
143+
private ArkResourceStsTokenInfo getEndpointToken(String endpointId, Integer ttl) throws ApiException {
144+
return getToken("endpoint", endpointId, ttl);
145+
}
146+
147+
private ArkResourceStsTokenInfo getToken(String resourceType, String resourceId, Integer ttl) throws ApiException {
148+
if (ttl < this.advisoryRefreshTimeout * 2) {
149+
throw new ArkException("ttl should not be under " + this.advisoryRefreshTimeout * 2 + " seconds.");
150+
}
151+
152+
GetApiKeyRequest r = new GetApiKeyRequest();
153+
r.durationSeconds(ttl);
154+
r.resourceType(resourceType);
155+
List<String> list = new ArrayList<>();
156+
list.add(resourceId);
157+
r.resourceIds(list);
158+
159+
GetApiKeyResponse response = this.volcClient.getApiKey(r);
160+
return new ArkResourceStsTokenInfo(response.getApiKey(), response.getExpiredTime());
161+
}
162+
163+
private String getResourceKey(String resourceType, String resourceId) {
164+
return resourceType + "/" + resourceId;
165+
}
166+
167+
public static class ArkResourceStsTokenInfo {
168+
private String token;
169+
private Integer expiredTime;
170+
171+
public ArkResourceStsTokenInfo(String token, Integer expiredTime) {
172+
this.token = token;
173+
this.expiredTime = expiredTime;
174+
}
175+
176+
public String getToken() {
177+
return token;
178+
}
179+
180+
public void setToken(String token) {
181+
this.token = token;
182+
}
183+
184+
public Integer getExpiredTime() {
185+
return expiredTime;
186+
}
187+
188+
public void setExpiredTime(Integer expiredTime) {
189+
this.expiredTime = expiredTime;
190+
}
191+
}
192+
}

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

Lines changed: 0 additions & 165 deletions
This file was deleted.

0 commit comments

Comments
 (0)