Skip to content

Commit 872a516

Browse files
zhou9584Le Zhou
zhou9584
and
Le Zhou
authored
[Feature] register agent with access token (#701)
<!-- Please provide brief information about the PR, what it contains & its purpose, new behaviors after the change. And let us know here if you need any help: https://github.com/microsoft/HydraLab/issues/new --> ## Description <!-- A few words to explain your changes --> ![image](https://github.com/user-attachments/assets/301c85f3-81e3-4fa4-9078-3fe649c2b4ab) ### Linked GitHub issue ID: # ## Pull Request Checklist <!-- Put an x in the boxes that apply. This is simply a reminder of what we are going to look for before merging your code. --> - [ ] Tests for the changes have been added (for bug fixes / features) - [x] Code compiles correctly with all tests are passed. - [x] I've read the [contributing guide](https://github.com/microsoft/HydraLab/blob/main/CONTRIBUTING.md#making-changes-to-the-code) and followed the recommended practices. - [ ] [Wikis](https://github.com/microsoft/HydraLab/wiki) or [README](https://github.com/microsoft/HydraLab/blob/main/README.md) have been reviewed and added / updated if needed (for bug fixes / features) ### Does this introduce a breaking change? *If this introduces a breaking change for Hydra Lab users, please describe the impact and migration path.* - [x] Yes - [ ] No ## How you tested it *Please make sure the change is tested, you can test it by adding UTs, do local test and share the screenshots, etc.* Please check the type of change your PR introduces: - [ ] Bugfix - [x] Feature - [ ] Technical design - [ ] Build related changes - [ ] Refactoring (no functional changes, no api changes) - [ ] Code style update (formatting, renaming) or Documentation content changes - [ ] Other (please describe): ### Feature UI screenshots or Technical design diagrams *If this is a relatively large or complex change, kick it off by drawing the tech design with PlantUML and explaining why you chose the solution you did and what alternatives you considered, etc...* --------- Co-authored-by: Le Zhou <[email protected]>
1 parent 7d41386 commit 872a516

File tree

13 files changed

+294
-17
lines changed

13 files changed

+294
-17
lines changed

‎agent/src/main/java/com/microsoft/hydralab/agent/service/AgentManageService.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33

44
package com.microsoft.hydralab.agent.service;
55

6+
import com.azure.core.credential.TokenCredential;
7+
import com.azure.core.credential.TokenRequestContext;
8+
import com.azure.identity.AzureCliCredentialBuilder;
69
import com.microsoft.hydralab.agent.config.AppOptions;
710
import com.microsoft.hydralab.common.entity.common.AgentUpdateTask;
811
import com.microsoft.hydralab.common.entity.common.Message;
@@ -93,4 +96,16 @@ public void restartAgent(String packageFileName, String path) {
9396
path);
9497
}
9598
}
99+
100+
public String getAgentAccessToken(String appId) {
101+
try {
102+
TokenCredential defaultAzureCredential = new AzureCliCredentialBuilder().build();
103+
TokenRequestContext tokenRequestContext = new TokenRequestContext().addScopes("api://" + appId);
104+
logger.info("Requesting access token for appId: {}", appId);
105+
return defaultAzureCredential.getToken(tokenRequestContext).block().getToken();
106+
} catch (Exception e) {
107+
logger.error("Failed to get agent access token for appId: {}", appId, e);
108+
return null;
109+
}
110+
}
96111
}

‎agent/src/main/java/com/microsoft/hydralab/agent/service/AgentWebSocketClientService.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,9 @@ public void onMessage(Message message) {
9292
Message response = null;
9393
switch (path) {
9494
case Const.Path.AUTH:
95+
if (!(message.getBody() instanceof JSONObject)) {
96+
break;
97+
}
9598
provideAuthInfo(message);
9699
return;
97100
case Const.Path.AGENT_INIT:
@@ -240,6 +243,8 @@ private void prometheusPushgatewayInit(AgentMetadata agentMetadata) {
240243
}
241244

242245
private void provideAuthInfo(Message message) {
246+
JSONObject authData = (JSONObject) message.getBody();
247+
String authMode = authData.getString(Const.AgentConfig.AGENT_AUTH_MODE_PARAM);
243248
Message responseAuth = new Message();
244249
responseAuth.setSessionId(message.getSessionId());
245250

@@ -248,7 +253,22 @@ private void provideAuthInfo(Message message) {
248253
}
249254
agentUser.setId(agentId);
250255
agentUser.setName(agentName);
251-
agentUser.setSecret(agentSecret);
256+
switch (authMode) {
257+
case Const.AgentAuthMode.TOKEN:
258+
String appId = authData.getString(Const.AgentConfig.AUTH_APP_ID_PARAM);
259+
agentUser.setSecret(agentManageService.getAgentAccessToken(appId));
260+
break;
261+
case Const.AgentAuthMode.SECRET:
262+
agentUser.setSecret(agentSecret);
263+
break;
264+
default:
265+
log.error("Unknown auth mode: {}", authMode);
266+
responseAuth.setCode(400);
267+
responseAuth.setMessage("Unknown auth mode: " + authMode);
268+
send(responseAuth);
269+
return;
270+
}
271+
252272
try {
253273
InetAddress localHost = InetAddress.getLocalHost();
254274
agentUser.setHostname(localHost.getHostName());

‎center/src/main/java/com/microsoft/hydralab/center/controller/AgentManageController.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,4 +295,27 @@ public Result deleteAgent(@CurrentSecurityContext SysUser requestor,
295295
agentManageService.deleteAgent(agentManageService.getAgent(agentId));
296296
return Result.ok("Delete Success");
297297
}
298+
299+
/**
300+
* Authenticated USER:
301+
* 1) users with ROLE SUPER_ADMIN/ADMIN,
302+
* 2) agent creator,
303+
* 3) admin of the TEAM that agent is in
304+
*/
305+
@PostMapping("/api/agent/updateDeviceId")
306+
public Result updateAgentDeviceId(@CurrentSecurityContext SysUser requestor,
307+
@RequestParam(value = "agentId", required = true) String agentId,
308+
@RequestParam(value = "deviceId", required = true) String deviceId) {
309+
if (!agentManageService.checkAgentAuthorization(requestor, agentId)) {
310+
return Result.error(HttpStatus.UNAUTHORIZED.value(), "Authentication failed");
311+
}
312+
313+
try {
314+
agentManageService.updateAgentDeviceId(agentId, deviceId);
315+
} catch (Exception e) {
316+
log.error("Failed to update device info for agent {}: {}", agentId, e.getMessage());
317+
return Result.error(HttpStatus.INTERNAL_SERVER_ERROR.value(), "Failed to update device info: " + e.getMessage());
318+
}
319+
return Result.ok("Update Device Info Success!");
320+
}
298321
}

‎center/src/main/java/com/microsoft/hydralab/center/interceptor/BaseInterceptor.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,6 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
7474

7575
// For Azure AD authentication
7676
String accessToken = request.getHeader("X-MS-TOKEN-AAD-ID-TOKEN");
77-
78-
LOGGER.info("IdToken: " + request.getHeader("X-MS-TOKEN-AAD-ID-TOKEN"));
79-
LOGGER.info("AccessToken: " + request.getHeader("X-MS-TOKEN-AAD-ACCESS-TOKEN"));
8077
LOGGER.info("UserId: " + request.getHeader("X-MS-CLIENT-PRINCIPAL-ID"));
8178
LOGGER.info("UserName: " + request.getHeader("X-MS-CLIENT-PRINCIPAL-NAME"));
8279

‎center/src/main/java/com/microsoft/hydralab/center/repository/AgentUserRepository.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,4 +24,6 @@ public interface AgentUserRepository extends JpaRepository<AgentUser, String>, J
2424
List<AgentUser> findAllByMailAddress(String mailAddress);
2525

2626
List<AgentUser> findAllByTeamId(String teamId);
27+
28+
Optional<AgentUser> findByDeviceId(String deviceId);
2729
}

‎center/src/main/java/com/microsoft/hydralab/center/service/AgentManageService.java‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,11 @@
99
import com.microsoft.hydralab.common.entity.center.SysUser;
1010
import com.microsoft.hydralab.common.entity.common.AgentUser;
1111
import com.microsoft.hydralab.common.entity.common.CriteriaType;
12+
import com.microsoft.hydralab.common.util.Const;
1213
import com.microsoft.hydralab.common.util.CriteriaTypeUtil;
14+
import com.microsoft.hydralab.common.util.HydraLabRuntimeException;
1315
import lombok.extern.slf4j.Slf4j;
16+
import org.springframework.beans.factory.annotation.Value;
1417
import org.springframework.data.jpa.domain.Specification;
1518
import org.springframework.stereotype.Component;
1619

@@ -30,6 +33,8 @@ public class AgentManageService {
3033
private UserTeamManagementService userTeamManagementService;
3134
@Resource
3235
private SysUserService sysUserService;
36+
@Value("${app.agent-auth-mode: 'SECRET'}")
37+
private String agentAuthMode;
3338

3439
public AgentUser createAgent(String teamId, String teamName, String mailAddress, String os, String name) {
3540
AgentUser agentUserInfo = new AgentUser();
@@ -51,7 +56,16 @@ public AgentUser createAgent(String teamId, String teamName, String mailAddress,
5156

5257
public AgentUser getAgent(String agentId) {
5358
Optional<AgentUser> agent = agentUserRepository.findById(agentId);
54-
return agent.orElse(null);
59+
if (!agent.isPresent()) {
60+
return null;
61+
}
62+
AgentUser agentUser = agent.get();
63+
// If agentAuthMode is TOKEN
64+
if (Const.AgentAuthMode.TOKEN.equals(agentAuthMode)) {
65+
// Set secret to "Not_Required" for TOKEN auth mode
66+
agentUser.setSecret("Not_Required");
67+
}
68+
return agentUser;
5569
}
5670

5771
public void deleteAgent(AgentUser agentUser) {
@@ -132,6 +146,16 @@ public void updateAgentTeam(String teamId, String teamName) {
132146
agentUserRepository.saveAll(agents);
133147
}
134148

149+
public void updateAgentDeviceId(String agentId, String deviceId) {
150+
Optional<AgentUser> findUser = agentUserRepository.findById(agentId);
151+
if (!findUser.isPresent()) {
152+
throw new HydraLabRuntimeException("Agent with ID " + agentId + " not found.");
153+
}
154+
AgentUser user = findUser.get();
155+
user.setDeviceId(deviceId);
156+
agentUserRepository.saveAndFlush(user);
157+
}
158+
135159
public File generateAgentConfigFile(String agentId, String host) {
136160
AgentUser agentUser = getAgent(agentId);
137161
if (agentUser != null) {
@@ -148,6 +172,9 @@ public File generateAgentConfigFile(String agentId, String host) {
148172
new File(CenterConstant.CENTER_TEMP_FILE_DIR));
149173

150174
FileWriter fileWriter = new FileWriter(agentConfigFile.getAbsolutePath());
175+
if(Const.AgentAuthMode.TOKEN.equals(agentAuthMode)){
176+
agentUser.setSecret("Not_Required");
177+
}
151178
fileWriter.write("app:\n" +
152179
" # register to Hydra Lab Center\n" +
153180
" registry:\n" +

‎center/src/main/java/com/microsoft/hydralab/center/service/DeviceAgentManagementService.java‎

Lines changed: 58 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.android.ddmlib.IDevice;
99
import com.microsoft.hydralab.center.openai.SuggestionService;
1010
import com.microsoft.hydralab.center.repository.AgentUserRepository;
11+
import com.microsoft.hydralab.center.util.AuthUtil;
1112
import com.microsoft.hydralab.center.util.MetricUtil;
1213
import com.microsoft.hydralab.common.entity.agent.AgentFunctionAvailability;
1314
import com.microsoft.hydralab.common.entity.agent.EnvCapabilityRequirement;
@@ -104,6 +105,8 @@ public class DeviceAgentManagementService {
104105
// save blocked devices <deviceSerial, blockedDeviceInfo>
105106
private final ConcurrentHashMap<String, BlockedDeviceInfo> blockedDevicesMap = new ConcurrentHashMap<>();
106107

108+
@Resource
109+
AuthUtil authUtil;
107110
@Resource
108111
BlockedDeviceInfoRepository blockedDeviceInfoRepository;
109112
@Resource
@@ -133,7 +136,6 @@ public class DeviceAgentManagementService {
133136

134137
@Value("${app.storage.type}")
135138
private String storageType;
136-
137139
@Value("${app.access-token-limit}")
138140
int accessLimit;
139141
@Value("${app.batteryStrategy}")
@@ -147,6 +149,8 @@ public class DeviceAgentManagementService {
147149
private boolean appCenterEnabled;
148150
@Value("${app.error-reporter.app-center.agent.secret: ''}")
149151
private String appCenterSecret;
152+
@Value("${app.agent-auth-mode: 'SECRET'}")
153+
private String agentAuthMode;
150154

151155
public void onOpen(Session session) {
152156
onlineCount.incrementAndGet();
@@ -193,7 +197,12 @@ private void sendAgentMetadata(Session session, AgentUser agentUser, String sign
193197
}
194198

195199
private void requestAuth(Session session) {
196-
sendMessageToSession(session, Message.auth());
200+
Message message = Message.auth();
201+
JSONObject authData = new JSONObject();
202+
authData.put(Const.AgentConfig.AGENT_AUTH_MODE_PARAM, agentAuthMode);
203+
authData.put(Const.AgentConfig.AUTH_APP_ID_PARAM, authUtil.getClientId());
204+
message.setBody(authData);
205+
sendMessageToSession(session, message);
197206
}
198207

199208
public void onMessage(Message message, @NotNull Session session) throws IOException {
@@ -617,16 +626,23 @@ private AgentUser searchQualifiedAgent(Message message) {
617626
return null;
618627
}
619628
AgentUser agentUser = (AgentUser) body;
620-
String id = agentUser.getId();
621-
622-
Optional<AgentUser> findUser = agentUserRepository.findById(id);
623-
if (!findUser.isPresent()) {
624-
return null;
629+
AgentUser user;
630+
switch (agentAuthMode) {
631+
case Const.AgentAuthMode.TOKEN:
632+
user = searchAgentByToken(agentUser);
633+
break;
634+
case Const.AgentAuthMode.SECRET:
635+
user = searchAgentBySecret(agentUser);
636+
break;
637+
default:
638+
log.error("Unknown agent auth mode: {}, please check your configuration.", agentAuthMode);
639+
return null;
625640
}
626-
AgentUser user = findUser.get();
627-
if (!user.getSecret().equals(agentUser.getSecret()) || !user.getName().equals(agentUser.getName())) {
641+
if (user == null) {
642+
log.warn("Agent user {} is not found or invalid.", agentUser);
628643
return null;
629644
}
645+
630646
user.setHostname(agentUser.getHostname());
631647
user.setIp(agentUser.getIp());
632648
user.setVersionName(agentUser.getVersionName());
@@ -636,6 +652,39 @@ private AgentUser searchQualifiedAgent(Message message) {
636652
return user;
637653
}
638654

655+
private AgentUser searchAgentByToken(AgentUser agentUser) {
656+
String accessToken = agentUser.getSecret();
657+
if (accessToken == null || accessToken.isEmpty()) {
658+
return null;
659+
}
660+
if (!authUtil.isValidToken(accessToken)) {
661+
return null;
662+
}
663+
String deviceId = authUtil.decodeAccessToken(accessToken).getString("deviceid");
664+
if (deviceId == null || deviceId.isEmpty()) {
665+
return null;
666+
}
667+
Optional<AgentUser> findUser = agentUserRepository.findByDeviceId(deviceId);
668+
if (!findUser.isPresent()) {
669+
return null;
670+
}
671+
return findUser.get();
672+
}
673+
674+
private AgentUser searchAgentBySecret(AgentUser agentUser) {
675+
String id = agentUser.getId();
676+
677+
Optional<AgentUser> findUser = agentUserRepository.findById(id);
678+
if (!findUser.isPresent()) {
679+
return null;
680+
}
681+
AgentUser user = findUser.get();
682+
if (!user.getSecret().equals(agentUser.getSecret())) {
683+
return null;
684+
}
685+
return user;
686+
}
687+
639688
public void onClose(Session session) {
640689
onlineCount.decrementAndGet();
641690
log.info("Session closed, id：{}，online count: {}", session.getId(), onlineCount.get());

0 commit comments

Comments
 (0)