/**
* 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.zookeeper.server.util;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.ConnectException;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* A utility that does bi-directional forwarding between two ports.
* Useful, for example, to simulate network failures.
* Example:
*
* Server 1 config file:
*
* server.1=127.0.0.1:7301:7401;8201
* server.2=127.0.0.1:7302:7402;8202
* server.3=127.0.0.1:7303:7403;8203
*
* Server 2 and 3 config files:
*
* server.1=127.0.0.1:8301:8401;8201
* server.2=127.0.0.1:8302:8402;8202
* server.3=127.0.0.1:8303:8403;8203
*
* Initially forward traffic between 730x and 830x and between 740x and 830x
* This way server 1 can communicate with servers 2 and 3
* ....
*
* List<PortForwarder> pfs = startForwarding();
* ....
* // simulate a network interruption for server 1
* stopForwarding(pfs);
* ....
* // restore connection
* pfs = startForwarding();
*
*
* private List<PortForwarder> startForwarding() throws IOException {
* List<PortForwarder> res = new ArrayList<PortForwarder>();
* res.add(new PortForwarder(8301, 7301));
* res.add(new PortForwarder(8401, 7401));
* res.add(new PortForwarder(7302, 8302));
* res.add(new PortForwarder(7402, 8402));
* res.add(new PortForwarder(7303, 8303));
* res.add(new PortForwarder(7403, 8403));
* return res;
* }
*
* private void stopForwarding(List<PortForwarder> pfs) throws Exception {
* for (PortForwarder pf : pfs) {
* pf.shutdown();
* }
* }
*
*
*/
public class PortForwarder extends Thread {
private static final Logger LOG = LoggerFactory
.getLogger(PortForwarder.class);
private static class PortForwardWorker implements Runnable {
private final InputStream in;
private final OutputStream out;
private final Socket toClose;
private final Socket toClose2;
PortForwardWorker(Socket toClose, Socket toClose2, InputStream in,
OutputStream out) throws IOException {
this.toClose = toClose;
this.toClose2 = toClose2;
this.in = in;
this.out = out;
// LOG.info("starting forward for "+toClose);
}
public void run() {
Thread.currentThread().setName(toClose.toString() + "-->"
+ toClose2.toString());
byte[] buf = new byte[1024];
try {
while (true) {
try {
int read = this.in.read(buf);
if (read > 0) {
try {
this.out.write(buf, 0, read);
} catch (IOException e) {
LOG.warn("exception during write", e);
try {
toClose.close();
} catch (IOException ex) {
// ignore
}
try {
toClose2.close();
} catch (IOException ex) {
// ignore
}
break;
}
}
} catch (SocketTimeoutException e) {
LOG.error("socket timeout", e);
}
Thread.sleep(1);
}
} catch (InterruptedException e) {
LOG.warn("Interrupted", e);
try {
toClose.close();
} catch (IOException ex) {
// ignore
}
try {
toClose2.close();
} catch (IOException ex) {
// ignore silently
}
} catch (SocketException e) {
if (!"Socket closed".equals(e.getMessage())) {
LOG.error("Unexpected exception", e);
}
} catch (IOException e) {
LOG.error("Unexpected exception", e);
}
LOG.info("Shutting down forward for " + toClose);
}
}
private volatile boolean stopped = false;
private ExecutorService workers = Executors.newCachedThreadPool();
private ServerSocket serverSocket;
private final int to;
public PortForwarder(int from, int to) throws IOException {
this.to = to;
serverSocket = new ServerSocket(from);
serverSocket.setSoTimeout(30000);
this.start();
}
@Override
public void run() {
try {
while (!stopped) {
Socket sock = null;
try {
LOG.info("accepting socket local:"
+ serverSocket.getLocalPort() + " to:" + to);
sock = serverSocket.accept();
LOG.info("accepted: local:" + sock.getLocalPort()
+ " from:" + sock.getPort()
+ " to:" + to);
Socket target = null;
int retry = 10;
while(sock.isConnected()) {
try {
target = new Socket("localhost", to);
break;
} catch (IOException e) {
if (retry == 0) {
throw e;
}
LOG.warn("connection failed, retrying(" + retry
+ "): local:" + sock.getLocalPort()
+ " from:" + sock.getPort()
+ " to:" + to, e);
}
Thread.sleep(TimeUnit.SECONDS.toMillis(1));
retry--;
}
LOG.info("connected: local:" + sock.getLocalPort()
+ " from:" + sock.getPort()
+ " to:" + to);
sock.setSoTimeout(30000);
target.setSoTimeout(30000);
this.workers.execute(new PortForwardWorker(sock, target,
sock.getInputStream(), target.getOutputStream()));
this.workers.execute(new PortForwardWorker(target, sock,
target.getInputStream(), sock.getOutputStream()));
} catch (SocketTimeoutException e) {
LOG.warn("socket timed out local:"
+ (sock != null ? sock.getLocalPort(): "")
+ " from:" + (sock != null ? sock.getPort(): "")
+ " to:" + to, e);
} catch (ConnectException e) {
LOG.warn("connection exception local:"
+ (sock != null ? sock.getLocalPort(): "")
+ " from:" + (sock != null ? sock.getPort(): "")
+ " to:" + to, e);
sock.close();
} catch (IOException e) {
if (!"Socket closed".equals(e.getMessage())) {
LOG.warn("unexpected exception local:"
+ (sock != null ? sock.getLocalPort(): "")
+ " from:" + (sock != null ? sock.getPort(): "")
+ " to:" + to, e);
throw e;
}
}
}
} catch (IOException e) {
LOG.error("Unexpected exception to:" + to, e);
} catch (InterruptedException e) {
LOG.error("Interrupted to:" + to, e);
}
}
public void shutdown() throws Exception {
this.stopped = true;
this.serverSocket.close();
this.workers.shutdownNow();
try {
if (!this.workers.awaitTermination(5, TimeUnit.SECONDS)) {
throw new Exception(
"Failed to stop forwarding within 5 seconds");
}
} catch (InterruptedException e) {
throw new Exception("Failed to stop forwarding");
}
this.join();
}
}