[Java] Custom headers/properties through Arrow Flight SQL JDBC

[aiguofer]

Describe the usage question you have. Please include as many useful details as possible.

We're currently building a Flight SQL Server and want to use CallHeaderAuthenticator. We want to use token based auth, but the user needs to provide some extra fields (accountId and environmentId). We need to validate the token, and then ensure the user is authorized for the provided accountId and environmentId.

When I pass custom properties to the JDBC URL, they appear in the headers server-side on some requests, but not all.

Here's what we have (CallHeadersWrapper simply creates an object with our custom headers out of passed in headers or a call context).

class CustomBearerTokenAuthenticator(
    private val authClient: AuthClient,
    private val logger: Logger
) :
    CallHeaderAuthenticator {
    private val authException = CallStatus.UNAUTHENTICATED.withDescription("Invalid credentials")

    override fun authenticate(incomingHeaders: CallHeaders?): CallHeaderAuthenticator.AuthResult {
        val bearerToken = AuthUtilities.getValueFromAuthHeader(incomingHeaders, Auth2Constants.BEARER_PREFIX)
            ?: throw authException.toRuntimeException()

        logger.info("Calling auth with headers: ${incomingHeaders.toString()}")
        val customHeaders: CallHeadersWrapper =
            incomingHeaders?.let { CallHeadersWrapper(it, logger) } ?: throw authException.toRuntimeException()

//        if (customHeaders.accountId == null) {
//            throw authException.withDescription("Parameter accountId is required").toRuntimeException()
//        }
//        if (customHeaders.environmentId == null) {
//            throw authException.withDescription("Parameter environmentId is required").toRuntimeException()
//        }
        // FIXME: looks like we don't have access to our custom headers during initial auth... some requests include our custom headers and some don't
        if (authClient.isValid(
                bearerToken,
                customHeaders.accountId ?: "1",
                customHeaders.environmentId ?: "12"
            )
        ) {
            return createAuthResult(bearerToken)
        } else {
            logger.info("Bearer token validation failed.")
            throw authException.toRuntimeException()
        }
    }

    private fun createAuthResult(token: String): CallHeaderAuthenticator.AuthResult {
        class AuthResult : CallHeaderAuthenticator.AuthResult {
            override fun appendToOutgoingHeaders(outgoingHeaders: CallHeaders?) {
                // TODO: we could store other info in headers to track session
                outgoingHeaders?.insert(
                    "authenticated", "yes"
                )
            }

            override fun getPeerIdentity(): String {
                // TODO: peer identity is an identifier for the actor, maybe we want to
                // include other things like accountId and projectId
                return "person"
            }
        }
        return AuthResult()
    }
}

When we use this JDBC connection string: jdbc:arrow-flight-sql://localhost:8811;useEncryption=false;accountId=1;envId=42;token=123 this is the output server side when I establish a connection to the server using DBeaver or DataGrip:

18:29:04.156 [flight-server-default-executor-3] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.161 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.165 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.170 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.173 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.178 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.180 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
18:29:04.183 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)

Am I misunderstanding something about how the headers are supposed to work here? any help would be greatly appreciated!

Component(s)

Java

[lidavidm]

@jduo do you happen to have an idea?

[jduo]

The intent is to always send custom properties with everything request. Let's try to breakdown what's happening with a standalone app that tries calling requests such as getDatabaseMetaData, runs queries, getTables, after connection

[lidavidm]

@aiguofer do you know which RPC call on the server is missing the headers? That'll help us track it down in the client

[aiguofer]

I'm not sure what is the first call the JDBC driver makes, but I can try to add some logging to find out. It looks like when I first create a connection it makes a bunch of requests, but I know for a fact the first one doesn't have the custom headers.

[aiguofer]

Looks like there's an initial call that is not getting the headers, then getFlightInfoSqlInfo gets them but getStreamSqlInfo does not.

10:50:42.165 [flight-server-default-executor-0] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
10:50:42.177 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
10:50:42.226 [flight-server-default-executor-1] INFO  logger -- getFlightInfoSqlInfo: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
10:50:42.371 [flight-server-default-executor-1] INFO  logger -- Calling auth with headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
10:50:42.373 [flight-server-default-executor-1] INFO  logger -- getStreamSqlInfo: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123

[lidavidm]

Thanks - IIRC gRPC interceptors get full access to the gRPC method name which would be helpful to narrow it down

[aiguofer]

Implemented my own CustomServerCallHeaderAuthMiddleware adding some logging:

Method: HANDSHAKE, headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)
Method: GET_FLIGHT_INFO, headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,envid=42,accountid=1,grpc-accept-encoding=gzip,authorization=Bearer 123)
Method: DO_GET, headers: Metadata(content-type=application/grpc,user-agent=grpc-java-netty/1.49.1,grpc-accept-encoding=gzip,authorization=Bearer 123)

[aiguofer]

The above PR should at least fix the issue for the HANDSHAKE. I'm not sure about the DO_GET, but for us that is not a blocker at this time.

[lidavidm]

Thanks for the PR. DoGet is probably missing a similar line somewhere. We can file a new issue once this first PR is merged.