/*
* 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.hive.spark.client.rpc;
import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.SecureRandom;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.TimeUnit;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.sasl.AuthorizeCallback;
import javax.security.sasl.RealmCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslException;
import javax.security.sasl.SaslServer;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.GenericFutureListener;
import io.netty.util.concurrent.Promise;
import io.netty.util.concurrent.ScheduledFuture;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.hadoop.hive.common.classification.InterfaceAudience;
/**
* An RPC server. The server matches remote clients based on a secret that is generated on
* the server - the secret needs to be given to the client through some other mechanism for
* this to work.
*/
@InterfaceAudience.Private
public class RpcServer implements Closeable {
private static final Logger LOG = LoggerFactory.getLogger(RpcServer.class);
private static final SecureRandom RND = new SecureRandom();
private final String address;
private final Channel channel;
private final EventLoopGroup group;
private final int port;
private final ConcurrentMap<String, ClientInfo> pendingClients;
private final RpcConfiguration config;
public RpcServer(Map<String, String> mapConf) throws IOException, InterruptedException {
this.config = new RpcConfiguration(mapConf);
this.group = new NioEventLoopGroup(
this.config.getRpcThreadCount(),
new ThreadFactoryBuilder()
.setNameFormat("RPC-Handler-%d")
.setDaemon(true)
.build());
ServerBootstrap serverBootstrap = new ServerBootstrap()
.group(group)
.channel(NioServerSocketChannel.class)
.childHandler(new ChannelInitializer<SocketChannel>() {
@Override
public void initChannel(SocketChannel ch) throws Exception {
SaslServerHandler saslHandler = new SaslServerHandler(config);
final Rpc newRpc = Rpc.createServer(saslHandler, config, ch, group);
saslHandler.rpc = newRpc;
Runnable cancelTask = new Runnable() {
@Override
public void run() {
LOG.warn("Timed out waiting for hello from client.");
newRpc.close();
}
};
saslHandler.cancelTask = group.schedule(cancelTask,
RpcServer.this.config.getConnectTimeoutMs(),
TimeUnit.MILLISECONDS);
}
})
.option(ChannelOption.SO_BACKLOG, 1)
.option(ChannelOption.SO_REUSEADDR, true)
.childOption(ChannelOption.SO_KEEPALIVE, true);
this.channel = bindServerPort(serverBootstrap).channel();
this.port = ((InetSocketAddress) channel.localAddress()).getPort();
this.pendingClients = Maps.newConcurrentMap();
this.address = this.config.getServerAddress();
}
/**
* Retry the list of configured ports until one is found
* @param serverBootstrap
* @return
* @throws InterruptedException
* @throws IOException
*/
private ChannelFuture bindServerPort(ServerBootstrap serverBootstrap)
throws InterruptedException, IOException {
List<Integer> ports = config.getServerPorts();
if (ports.contains(0)) {
return serverBootstrap.bind(0).sync();
} else {
Random rand = new Random();
while(!ports.isEmpty()) {
int index = rand.nextInt(ports.size());
int port = ports.get(index);
ports.remove(index);
try {
return serverBootstrap.bind(port).sync();
} catch(Exception e) {
// Retry the next port
}
}
throw new IOException("No available ports from configured RPC Server ports for HiveServer2");
}
}
/**
* Tells the RPC server to expect a connection from a new client.
*
* @param clientId An identifier for the client. Must be unique.
* @param secret The secret the client will send to the server to identify itself.
* @param serverDispatcher The dispatcher to use when setting up the RPC instance.
* @return A future that can be used to wait for the client connection, which also provides the
* secret needed for the client to connect.
*/
public Future<Rpc> registerClient(final String clientId, String secret,
RpcDispatcher serverDispatcher) {
return registerClient(clientId, secret, serverDispatcher, config.getServerConnectTimeoutMs());
}
@VisibleForTesting
Future<Rpc> registerClient(final String clientId, String secret,
RpcDispatcher serverDispatcher, long clientTimeoutMs) {
final Promise<Rpc> promise = group.next().newPromise();
Runnable timeout = new Runnable() {
@Override
public void run() {
promise.setFailure(new TimeoutException("Timed out waiting for client connection."));
}
};
ScheduledFuture<?> timeoutFuture = group.schedule(timeout,
clientTimeoutMs,
TimeUnit.MILLISECONDS);
final ClientInfo client = new ClientInfo(clientId, promise, secret, serverDispatcher,
timeoutFuture);
if (pendingClients.putIfAbsent(clientId, client) != null) {
throw new IllegalStateException(
String.format("Client '%s' already registered.", clientId));
}
promise.addListener(new GenericFutureListener<Promise<Rpc>>() {
@Override
public void operationComplete(Promise<Rpc> p) {
if (!p.isSuccess()) {
pendingClients.remove(clientId);
}
}
});
return promise;
}
/**
* Tells the RPC server to cancel the connection from an existing pending client
* @param clientId The identifier for the client
* @param msg The error message about why the connection should be canceled
*/
public void cancelClient(final String clientId, final String msg) {
final ClientInfo cinfo = pendingClients.remove(clientId);
if (cinfo == null) {
// Nothing to be done here.
return;
}
cinfo.timeoutFuture.cancel(true);
if (!cinfo.promise.isDone()) {
cinfo.promise.setFailure(new RuntimeException(
String.format("Cancel client '%s'. Error: " + msg, clientId)));
}
}
/**
* Creates a secret for identifying a client connection.
*/
public String createSecret() {
byte[] secret = new byte[config.getSecretBits() / 8];
RND.nextBytes(secret);
StringBuilder sb = new StringBuilder();
for (byte b : secret) {
if (b < 10) {
sb.append("0");
}
sb.append(Integer.toHexString(b));
}
return sb.toString();
}
public String getAddress() {
return address;
}
public int getPort() {
return port;
}
@Override
public void close() {
try {
channel.close();
for (ClientInfo client : pendingClients.values()) {
client.promise.cancel(true);
}
pendingClients.clear();
} finally {
group.shutdownGracefully();
}
}
private class SaslServerHandler extends SaslHandler implements CallbackHandler {
private final SaslServer server;
private Rpc rpc;
private ScheduledFuture<?> cancelTask;
private String clientId;
private ClientInfo client;
SaslServerHandler(RpcConfiguration config) throws IOException {
super(config);
this.server = Sasl.createSaslServer(config.getSaslMechanism(), Rpc.SASL_PROTOCOL,
Rpc.SASL_REALM, config.getSaslOptions(), this);
}
@Override
protected boolean isComplete() {
return server.isComplete();
}
@Override
protected String getNegotiatedProperty(String name) {
return (String) server.getNegotiatedProperty(name);
}
@Override
protected Rpc.SaslMessage update(Rpc.SaslMessage challenge) throws IOException {
if (clientId == null) {
Preconditions.checkArgument(challenge.clientId != null,
"Missing client ID in SASL handshake.");
clientId = challenge.clientId;
client = pendingClients.get(clientId);
Preconditions.checkArgument(client != null,
"Unexpected client ID '%s' in SASL handshake.", clientId);
}
return new Rpc.SaslMessage(server.evaluateResponse(challenge.payload));
}
@Override
public byte[] wrap(byte[] data, int offset, int len) throws IOException {
return server.wrap(data, offset, len);
}
@Override
public byte[] unwrap(byte[] data, int offset, int len) throws IOException {
return server.unwrap(data, offset, len);
}
@Override
public void dispose() throws IOException {
if (!server.isComplete()) {
onError(new SaslException("Server closed before SASL negotiation finished."));
}
server.dispose();
}
@Override
protected void onComplete() throws Exception {
cancelTask.cancel(true);
client.timeoutFuture.cancel(true);
rpc.setDispatcher(client.dispatcher);
client.promise.setSuccess(rpc);
pendingClients.remove(client.id);
}
@Override
protected void onError(Throwable error) {
cancelTask.cancel(true);
if (client != null) {
client.timeoutFuture.cancel(true);
if (!client.promise.isDone()) {
client.promise.setFailure(error);
}
}
}
@Override
public void handle(Callback[] callbacks) {
Preconditions.checkState(client != null, "Handshake not initialized yet.");
for (Callback cb : callbacks) {
if (cb instanceof NameCallback) {
((NameCallback)cb).setName(clientId);
} else if (cb instanceof PasswordCallback) {
((PasswordCallback)cb).setPassword(client.secret.toCharArray());
} else if (cb instanceof AuthorizeCallback) {
((AuthorizeCallback) cb).setAuthorized(true);
} else if (cb instanceof RealmCallback) {
RealmCallback rb = (RealmCallback) cb;
rb.setText(rb.getDefaultText());
}
}
}
}
private static class ClientInfo {
final String id;
final Promise<Rpc> promise;
final String secret;
final RpcDispatcher dispatcher;
final ScheduledFuture<?> timeoutFuture;
private ClientInfo(String id, Promise<Rpc> promise, String secret, RpcDispatcher dispatcher,
ScheduledFuture<?> timeoutFuture) {
this.id = id;
this.promise = promise;
this.secret = secret;
this.dispatcher = dispatcher;
this.timeoutFuture = timeoutFuture;
}
}
}