Skip to content

Commit

Permalink
7
Browse files Browse the repository at this point in the history
  • Loading branch information
xinyiZzz committed Sep 15, 2023
1 parent f8ef6df commit 3524285
Show file tree
Hide file tree
Showing 7 changed files with 257 additions and 4 deletions.
9 changes: 9 additions & 0 deletions fe/fe-core/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -768,6 +768,15 @@ under the License.
<groupId>org.apache.arrow</groupId>
<artifactId>flight-sql</artifactId>
</dependency>

<dependency>
<groupId>org.immutables</groupId>
<artifactId>value</artifactId>
</dependency>
<dependency>
<groupId>org.immutables</groupId>
<artifactId>value-annotations</artifactId>
</dependency>
</dependencies>
<repositories>
<!-- for huawei obs sdk -->
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
// This file is copied from
// https://github.com/dremio/dremio-oss/blob/master/services/arrow-flight/src/main/java/com/dremio/service/flight/ServerCookieMiddleware.java
// and modified by Doris

package org.apache.doris.service.arrowflight;

import java.util.Date;

import org.immutables.value.Value;

/**
* Result of Authentication.
*/
@Value.Immutable
public interface AuthResult {
String getUserName();
Date getExpiresAt();
String getUserId();

static AuthResult of(String userName) {
return new ImmutableAuthResult.Builder()
.setUserName(userName)
.build();
}

static ImmutableAuthResult.Builder builder() {
return new ImmutableAuthResult.Builder();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,10 @@

import com.google.common.base.Charsets;
import org.apache.arrow.flight.auth.BasicServerAuthHandler;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.Arrays;
import java.util.Optional;

/**
Expand All @@ -31,6 +34,7 @@
* TODO
*/
public class FlightServerBasicAuthValidator implements BasicServerAuthHandler.BasicAuthValidator {
private static final Logger LOG = LogManager.getLogger(FlightServerBasicAuthValidator.class);

private static final String myToken = "DORIS_READ_WRITE_TOKEN";

Expand All @@ -39,12 +43,14 @@ public FlightServerBasicAuthValidator() {

@Override
public byte[] getToken(String username, String password) {
LOG.info("getToken " + username + password);
return myToken.getBytes(Charsets.UTF_8);
}

@Override
public Optional<String> isValid(byte[] bytes) {
// final String token = new String(bytes, Charsets.UTF_8);
LOG.info("isValid " + Arrays.toString(bytes));
return Optional.of(myToken);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
// This file is copied from
// https://github.com/dremio/dremio-oss/blob/master/services/arrow-flight/src/main/java/com/dremio/service/flight/ServerCookieMiddleware.java
// and modified by Doris

package org.apache.doris.service.arrowflight;

import com.google.common.annotations.VisibleForTesting;
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.auth2.Auth2Constants;
import org.apache.arrow.flight.auth2.AuthUtilities;
import org.apache.arrow.flight.auth2.BasicCallHeaderAuthenticator;
import org.apache.arrow.flight.auth2.CallHeaderAuthenticator;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class FlightServerBearerTokenAuthenticator implements CallHeaderAuthenticator {
private static final Logger LOG = LogManager.getLogger(FlightServerBearerTokenAuthenticator.class);

private final CallHeaderAuthenticator initialAuthenticator;

public FlightServerBearerTokenAuthenticator() {
this.initialAuthenticator = new BasicCallHeaderAuthenticator(new FlightServerCredentialValidator());
}

@Override
public AuthResult authenticate(CallHeaders incomingHeaders) {
final String bearerToken = AuthUtilities.getValueFromAuthHeader(incomingHeaders,
Auth2Constants.BEARER_PREFIX);

if (bearerToken != null) {
return validateBearer(bearerToken);
} else {
final AuthResult result = initialAuthenticator.authenticate(incomingHeaders);
return getAuthResultWithBearerToken(result, incomingHeaders);
}
}

/**
* Validates provided token.
*
* @param token the token to validate.
* @return an AuthResult with the bearer token and peer identity.
*/
@VisibleForTesting
AuthResult validateBearer(String token) {
try {
// tokenManagerProvider.get().validateToken(token);
return createAuthResultWithBearerToken(token);
} catch (IllegalArgumentException e) {
LOG.error("Bearer token validation failed.", e);
throw CallStatus.UNAUTHENTICATED.toRuntimeException();
}
}

/**
* Generates a bearer token, parses client properties from incoming headers, then creates a
* UserSession associated with the generated token and client properties.
*
* @param authResult the AuthResult from initial authentication, with peer identity captured.
* @param incomingHeaders the CallHeaders to parse client properties from.
* @return an an AuthResult with the bearer token and peer identity.
*/
@VisibleForTesting
AuthResult getAuthResultWithBearerToken(AuthResult authResult, CallHeaders incomingHeaders) {
final String username = authResult.getPeerIdentity();
// final String token = DremioFlightAuthUtils.createUserSessionWithTokenAndProperties(
// tokenManagerProvider,
// username);

return createAuthResultWithBearerToken(username);
}

/**
* Helper method to create an AuthResult.
*
* @param token the token to create a UserSession for.
* @return a new AuthResult with functionality to add given bearer token to the outgoing header.
*/
private AuthResult createAuthResultWithBearerToken(String token) {
return new AuthResult() {
@Override
public void appendToOutgoingHeaders(CallHeaders outgoingHeaders) {
outgoingHeaders.insert(Auth2Constants.AUTHORIZATION_HEADER,
Auth2Constants.BEARER_PREFIX + token);
}

@Override
public String getPeerIdentity() {
return token;
}
};
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
// This file is copied from
// https://github.com/dremio/dremio-oss/blob/master/services/arrow-flight/src/main/java/com/dremio/service/flight/ServerCookieMiddleware.java
// and modified by Doris

package org.apache.doris.service.arrowflight;

import org.apache.doris.analysis.UserIdentity;
import org.apache.doris.catalog.Env;
import org.apache.doris.common.AuthenticationException;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.auth2.BasicCallHeaderAuthenticator;
import org.apache.arrow.flight.auth2.CallHeaderAuthenticator.AuthResult;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.List;

/**
* Dremio authentication specialized CredentialValidator implementation.
*/
public class FlightServerCredentialValidator implements BasicCallHeaderAuthenticator.CredentialValidator {
private static final Logger LOG = LogManager.getLogger(FlightServerCredentialValidator.class);

/**
* Authenticates against Dremio with the provided username and password.
*
* @param username Dremio username.
* @param password Dremio user password.
* @return AuthResult with username as the peer identity.
*/
@Override
public AuthResult validate(String username, String password) {
// String remoteIp = context.getMysqlChannel().getRemoteIp();
String remoteIp = "0.0.0.0";
List<UserIdentity> currentUserIdentity = Lists.newArrayList();

try {
Env.getCurrentEnv().getAuth().checkPlainPassword(username, remoteIp, password, currentUserIdentity);
} catch (AuthenticationException e) {
LOG.error("Unable to authenticate user {}", username, e);
final String errorMessage = "Unable to authenticate user " + username + ", exception: " + e.getMessage();
throw CallStatus.UNAUTHENTICATED.withCause(e).withDescription(errorMessage).toRuntimeException();
}
Preconditions.checkState(currentUserIdentity.size() == 1);
// context.setCurrentUserIdentity(currentUserIdentity.get(0));
// context.setRemoteIP(remoteIp);
org.apache.doris.service.arrowflight.AuthResult authResult = org.apache.doris.service.arrowflight.AuthResult.of(
username);
return authResult::getUserName;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightServerMiddleware;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.auth.BasicServerAuthHandler;
// import org.apache.arrow.flight.auth.BasicServerAuthHandler;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.logging.log4j.LogManager;
Expand All @@ -44,9 +44,10 @@ public FlightSqlService(int port) {
Location location = Location.forGrpcInsecure("0.0.0.0", port);
FlightSqlServiceImpl producer = new FlightSqlServiceImpl(location);
flightServer = FlightServer.builder(allocator, location, producer)
.middleware(FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY,
new FlightServerCookieMiddleware.Factory())
.authHandler(new BasicServerAuthHandler(new FlightServerBasicAuthValidator())).build();
.headerAuthenticator(new FlightServerBearerTokenAuthenticator()).build();
// .middleware(FLIGHT_CLIENT_PROPERTIES_MIDDLEWARE_KEY,
// new FlightServerCookieMiddleware.Factory())
// .authHandler(new BasicServerAuthHandler(new FlightServerBasicAuthValidator())).build();
}

// start Arrow Flight SQL service, return true if success, otherwise false
Expand Down
12 changes: 12 additions & 0 deletions fe/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,7 @@ under the License.
<woodstox.version>6.5.1</woodstox.version>
<kerby.version>2.0.3</kerby.version>
<jettison.version>1.5.4</jettison.version>
<immutables.version>2.9.3</immutables.version>
<vesoft.client.version>3.0.0</vesoft.client.version>
<!-- paimon -->
<paimon.version>0.4.0-incubating</paimon.version>
Expand Down Expand Up @@ -1489,6 +1490,17 @@ under the License.
<artifactId>arrow-jdbc</artifactId>
<version>${arrow.version}</version>
</dependency>
<dependency>
<groupId>org.immutables</groupId>
<artifactId>value</artifactId>
<version>${immutables.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.immutables</groupId>
<artifactId>value-annotations</artifactId>
<version>${immutables.version}</version>
</dependency>
</dependencies>
</dependencyManagement>
<dependencies>
Expand Down

0 comments on commit 3524285

Please sign in to comment.