jduo commented on a change in pull request #8583: URL: https://github.com/apache/arrow/pull/8583#discussion_r517017968
########## File path: java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java ########## @@ -0,0 +1,80 @@ +/* + * 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. + */ + +package org.apache.arrow.flight; + +/** + * Middleware to add a SET-SESSION header. + */ +public class ServerSessionMiddleware implements FlightServerMiddleware { + + /** + * Factory class to provide instances of ServerSessionMiddleware for each call. + */ + public static class Factory implements FlightServerMiddleware.Factory<ServerSessionMiddleware> { + private final ServerSessionHandler sessionHandler; + + public Factory(ServerSessionHandler sessionHandler) { + this.sessionHandler = sessionHandler; + } + + @Override + public ServerSessionMiddleware onCallStarted(CallInfo info, CallHeaders incomingHeaders, + RequestContext context) { + if (incomingHeaders.containsKey(FlightConstants.SESSION_HEADER)) { + // Client is re-using pre-existing session ID. + // ServerSessionHandler validates client session ID before proceeding. + final String sessionId = incomingHeaders.get(FlightConstants.SESSION_HEADER); + if (!sessionHandler.isValid(sessionId)) { + throw CallStatus.UNAUTHENTICATED.toRuntimeException(); + } + } else { + // Insert SET-SESSION header if ServerSessionHandler returns a non-null session ID. + final String sessionId = sessionHandler.getSessionID(); Review comment: Let's rename getSessionID() to beginSession() and have it take in the input headers. Let's have ServerSessionMiddleware take in the session id. ########## File path: java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java ########## @@ -0,0 +1,81 @@ +/* + * 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. + */ + +package org.apache.arrow.flight; + +/** + * Middleware to add a SET-SESSION header. + */ +public class ServerSessionMiddleware implements FlightServerMiddleware { + + /** + * Factory class to provide instances of ServerSessionMiddleware for each call. + */ + public static class Factory implements FlightServerMiddleware.Factory<ServerSessionMiddleware> { + private final ServerSessionHandler sessionHandler; + + public Factory(ServerSessionHandler sessionHandler) { + this.sessionHandler = sessionHandler; + } + + @Override + public ServerSessionMiddleware onCallStarted(CallInfo info, CallHeaders incomingHeaders, + RequestContext context) { + if (incomingHeaders.containsKey(FlightConstants.SESSION_HEADER)) { + // Client is re-using existing session ID. + // ServerSessionHandler validates client session ID before proceeding. + final String existingSessionId = incomingHeaders.get(FlightConstants.SESSION_HEADER); + if (!sessionHandler.isValid(existingSessionId)) { + throw CallStatus.UNAUTHENTICATED.toRuntimeException(); + } + } else { + // No existing session ID provided, establishing a new session. + // Insert SET-SESSION header if ServerSessionHandler returns a non-null session ID. + final String sessionId = sessionHandler.getSessionId(); + if (sessionId != null) { + incomingHeaders.insert(FlightConstants.SESSION_HEADER, sessionId); + } + } + + return new ServerSessionMiddleware(sessionHandler); + } + } + + private ServerSessionHandler sessionHandler; + + private ServerSessionMiddleware(ServerSessionHandler sessionHandler) { + this.sessionHandler = sessionHandler; + } + + @Override + public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { + String sessionId = sessionHandler.getSessionId(); Review comment: This should not call getSessionId(), because you can't guarantee that this sessionId is different from the one from onCallStarted(). Also, getSessionId() might trigger the backend starting a new session, so you'd double-start it. ########## File path: java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionMiddleware.java ########## @@ -0,0 +1,81 @@ +/* + * 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. + */ + +package org.apache.arrow.flight; + +/** + * Middleware to add a SET-SESSION header. + */ +public class ServerSessionMiddleware implements FlightServerMiddleware { + + /** + * Factory class to provide instances of ServerSessionMiddleware for each call. + */ + public static class Factory implements FlightServerMiddleware.Factory<ServerSessionMiddleware> { + private final ServerSessionHandler sessionHandler; + + public Factory(ServerSessionHandler sessionHandler) { + this.sessionHandler = sessionHandler; + } + + @Override + public ServerSessionMiddleware onCallStarted(CallInfo info, CallHeaders incomingHeaders, + RequestContext context) { + if (incomingHeaders.containsKey(FlightConstants.SESSION_HEADER)) { + // Client is re-using existing session ID. + // ServerSessionHandler validates client session ID before proceeding. + final String existingSessionId = incomingHeaders.get(FlightConstants.SESSION_HEADER); + if (!sessionHandler.isValid(existingSessionId)) { + throw CallStatus.UNAUTHENTICATED.toRuntimeException(); + } + } else { + // No existing session ID provided, establishing a new session. + // Insert SET-SESSION header if ServerSessionHandler returns a non-null session ID. + final String sessionId = sessionHandler.getSessionId(); Review comment: Rename getSessionId() to beginSession, and cache it instead of the SessionHandler in ServerSessionMiddleware. ########## File path: java/flight/flight-core/src/main/java/org/apache/arrow/flight/ServerSessionHandler.java ########## @@ -0,0 +1,56 @@ +/* + * 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. + */ + +package org.apache.arrow.flight; + +/** + * ServerSessionHandler interface to retrieve the current session ID. + */ +public interface ServerSessionHandler { + /** + * A session handler that does not support sessions. + * It is used as the default handler for a Flight server with no session capabilities. + */ + ServerSessionHandler NO_OP = new ServerSessionHandler() { + @Override + public String getSessionId() { + return null; + } + + @Override + public boolean isValid(String sessionId) { + return false; + } + }; + + /** + * Retrieves the current session ID. Generates a new session ID if there is no pre-existing Review comment: We need a way to generate a session ID based off headers. There is no 'current' session because this ServerSessionHandler can be servicing multiple different sessions at once. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
