Skip to content

Commit

Permalink
[fix](audit-loader) fix invalid token check logic (#32095)
Browse files Browse the repository at this point in the history
The check of the token should be forwarded to Master FE.
I add a new RPC method `checkToken()` in Frontend for this logic.
Otherwise, after enable the audit loader, the log from non-master FE can not be loaded to audit table
with `Invalid token` error.
  • Loading branch information
morningman authored Mar 12, 2024
1 parent 23aaf0f commit 1aa2d91
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import org.apache.doris.common.Config;
import org.apache.doris.common.DdlException;
import org.apache.doris.common.LoadException;
import org.apache.doris.common.UserException;
import org.apache.doris.httpv2.entity.ResponseEntityBuilder;
import org.apache.doris.httpv2.entity.RestBaseResult;
import org.apache.doris.httpv2.exception.UnauthorizedException;
Expand Down Expand Up @@ -362,7 +363,11 @@ private TNetworkAddress selectRedirectBackend(boolean groupCommit) throws LoadEx
// temporarily addressing the users' needs for audit logs.
// So this function is not widely tested under general scenario
private boolean checkClusterToken(String token) {
return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
try {
return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
} catch (UserException e) {
throw new UnauthorizedException(e.getMessage());
}
}

// NOTE: This function can only be used for AuditlogPlugin stream load for now.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,11 +63,6 @@ private String generateNewToken() {
return UUID.randomUUID().toString();
}

// this method only will be called in master node, since stream load only send message to master.
public boolean checkAuthToken(String token) {
return tokenQueue.contains(token);
}

public String acquireToken() throws UserException {
if (Env.getCurrentEnv().isMaster() || FeConstants.runningUnitTest) {
return tokenQueue.peek();
Expand All @@ -81,9 +76,8 @@ public String acquireToken() throws UserException {
}
}

public String acquireTokenFromMaster() throws TException {
private String acquireTokenFromMaster() throws TException {
TNetworkAddress thriftAddress = getMasterAddress();

FrontendService.Client client = getClient(thriftAddress);

if (LOG.isDebugEnabled()) {
Expand All @@ -108,7 +102,7 @@ public String acquireTokenFromMaster() throws TException {
} else {
TMySqlLoadAcquireTokenResult result = client.acquireToken();
if (result.getStatus().getStatusCode() != TStatusCode.OK) {
throw new TException("commit failed.");
throw new TException("acquire token from master failed. " + result.getStatus());
}
isReturnToPool = true;
return result.getToken();
Expand All @@ -122,6 +116,57 @@ public String acquireTokenFromMaster() throws TException {
}
}

/**
* Check if the token is valid.
* If this is not Master FE, will send the request to Master FE.
*/
public boolean checkAuthToken(String token) throws UserException {
if (Env.getCurrentEnv().isMaster() || FeConstants.runningUnitTest) {
return tokenQueue.contains(token);
} else {
try {
return checkTokenFromMaster(token);
} catch (TException e) {
LOG.warn("check token error", e);
throw new UserException("Check token from master failed", e);
}
}
}

private boolean checkTokenFromMaster(String token) throws TException {
TNetworkAddress thriftAddress = getMasterAddress();
FrontendService.Client client = getClient(thriftAddress);

if (LOG.isDebugEnabled()) {
LOG.debug("Send check token to Master {}", thriftAddress);
}

boolean isReturnToPool = false;
try {
boolean result = client.checkToken(token);
isReturnToPool = true;
return result;
} catch (TTransportException e) {
boolean ok = ClientPool.frontendPool.reopen(client, thriftTimeoutMs);
if (!ok) {
throw e;
}
if (e.getType() == TTransportException.TIMED_OUT) {
throw e;
} else {
boolean result = client.checkToken(token);
isReturnToPool = true;
return result;
}
} finally {
if (isReturnToPool) {
ClientPool.frontendPool.returnObject(thriftAddress, client);
} else {
ClientPool.frontendPool.invalidateObject(thriftAddress, client);
}
}
}


private TNetworkAddress getMasterAddress() throws TException {
Env.getCurrentEnv().checkReadyOrThrowTException();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1062,12 +1062,6 @@ private void checkPasswordAndPrivs(String user, String passwd, String db, List<S
}
}

private void checkToken(String token) throws AuthenticationException {
if (!Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token)) {
throw new AuthenticationException("Un matched cluster token.");
}
}

private void checkPassword(String user, String passwd, String clientIp)
throws AuthenticationException {
final String fullUserName = ClusterNamespace.getNameFromFullName(user);
Expand Down Expand Up @@ -1132,7 +1126,9 @@ private TLoadTxnBeginResult loadTxnBeginImpl(TLoadTxnBeginRequest request, Strin
request.getTbl(),
request.getUserIp(), PrivPredicate.LOAD);
} else {
checkToken(request.getToken());
if (!checkToken(request.getToken())) {
throw new AuthenticationException("Invalid token: " + request.getToken());
}
}

// check label
Expand Down Expand Up @@ -1343,7 +1339,9 @@ private void loadTxnPreCommitImpl(TLoadTxnCommitRequest request) throws UserExce
if (request.isSetAuthCode()) {
// CHECKSTYLE IGNORE THIS LINE
} else if (request.isSetToken()) {
checkToken(request.getToken());
if (!checkToken(request.getToken())) {
throw new AuthenticationException("Invalid token: " + request.getToken());
}
} else {
// refactoring it
if (CollectionUtils.isNotEmpty(request.getTbls())) {
Expand Down Expand Up @@ -2413,6 +2411,20 @@ public TMySqlLoadAcquireTokenResult acquireToken() throws TException {
return result;
}

@Override
public boolean checkToken(String token) {
String clientAddr = getClientAddrAsString();
if (LOG.isDebugEnabled()) {
LOG.debug("receive check token request from client: {}", clientAddr);
}
try {
return Env.getCurrentEnv().getLoadManager().getTokenManager().checkAuthToken(token);
} catch (Throwable e) {
LOG.warn("catch unknown result.", e);
return false;
}
}

@Override
public TCheckAuthResult checkAuth(TCheckAuthRequest request) throws TException {
String clientAddr = getClientAddrAsString();
Expand Down
2 changes: 2 additions & 0 deletions gensrc/thrift/FrontendService.thrift
Original file line number Diff line number Diff line change
Expand Up @@ -1462,6 +1462,8 @@ service FrontendService {

TMySqlLoadAcquireTokenResult acquireToken()

bool checkToken(1: string token)

TConfirmUnusedRemoteFilesResult confirmUnusedRemoteFiles(1: TConfirmUnusedRemoteFilesRequest request)

TCheckAuthResult checkAuth(1: TCheckAuthRequest request)
Expand Down

0 comments on commit 1aa2d91

Please sign in to comment.