View Javadoc
1   /*
2    * ====================================================================
3    * Licensed to the Apache Software Foundation (ASF) under one
4    * or more contributor license agreements.  See the NOTICE file
5    * distributed with this work for additional information
6    * regarding copyright ownership.  The ASF licenses this file
7    * to you under the Apache License, Version 2.0 (the
8    * "License"); you may not use this file except in compliance
9    * with the License.  You may obtain a copy of the License at
10   *
11   *   http://www.apache.org/licenses/LICENSE-2.0
12   *
13   * Unless required by applicable law or agreed to in writing,
14   * software distributed under the License is distributed on an
15   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16   * KIND, either express or implied.  See the License for the
17   * specific language governing permissions and limitations
18   * under the License.
19   * ====================================================================
20   *
21   * This software consists of voluntary contributions made by many
22   * individuals on behalf of the Apache Software Foundation.  For more
23   * information on the Apache Software Foundation, please see
24   * <http://www.apache.org/>.
25   *
26   */
27  package org.apache.hc.core5.testing;
28  
29  import java.io.DataInputStream;
30  import java.io.DataOutputStream;
31  import java.io.IOException;
32  import java.io.InputStream;
33  import java.io.OutputStream;
34  import java.net.InetAddress;
35  import java.net.ServerSocket;
36  import java.net.Socket;
37  import java.net.SocketAddress;
38  import java.util.ArrayList;
39  import java.util.List;
40  
41  import org.apache.hc.core5.net.InetAddressUtils;
42  import org.apache.hc.core5.util.TimeValue;
43  
44  /**
45   * Cheap and nasty SOCKS protocol version 5 proxy, recommended for use in unit tests only so we can test our SOCKS client code.
46   */
47  public class SocksProxy {
48  
49      private static class SocksProxyHandler {
50  
51          public static final int VERSION_5 = 5;
52          public static final int COMMAND_CONNECT = 1;
53          public static final int ATYP_DOMAINNAME = 3;
54  
55          private final SocksProxy parent;
56          private final Socket socket;
57          private volatile Socket remote;
58  
59          public SocksProxyHandler(final SocksProxy parent, final Socket socket) {
60              this.parent = parent;
61              this.socket = socket;
62          }
63  
64          public void start() {
65              new Thread(new Runnable() {
66                  @Override
67                  public void run() {
68                      try {
69                          final DataInputStream input = new DataInputStream(socket.getInputStream());
70                          final DataOutputStream output = new DataOutputStream(socket.getOutputStream());
71                          final Socket target = establishConnection(input, output);
72                          remote = target;
73  
74                          final Thread t1 = pumpStream(input, target.getOutputStream());
75                          final Thread t2 = pumpStream(target.getInputStream(), output);
76                          try {
77                              t1.join();
78                          } catch (final InterruptedException e) {
79                          }
80                          try {
81                              t2.join();
82                          } catch (final InterruptedException e) {
83                          }
84                      } catch (final IOException e) {
85                      } finally {
86                          parent.cleanupSocksProxyHandler(SocksProxyHandler.this);
87                      }
88                  }
89  
90                  private Socket establishConnection(final DataInputStream input, final DataOutputStream output) throws IOException {
91                      final int clientVersion = input.readUnsignedByte();
92                      if (clientVersion != VERSION_5) {
93                          throw new IOException("SOCKS implementation only supports version 5");
94                      }
95                      final int nMethods = input.readUnsignedByte();
96                      for (int i = 0; i < nMethods; i++) {
97                          input.readUnsignedByte(); // auth method
98                      }
99                      // response
100                     output.writeByte(VERSION_5);
101                     output.writeByte(0); // no auth method
102                     output.flush();
103 
104                     input.readUnsignedByte(); // client version again
105                     final int command = input.readUnsignedByte();
106                     if (command != COMMAND_CONNECT) {
107                         throw new IOException("SOCKS implementation only supports CONNECT command");
108                     }
109                     input.readUnsignedByte(); // reserved
110 
111                     final String targetHost;
112                     final byte[] targetAddress;
113                     final int addressType = input.readUnsignedByte();
114                     switch (addressType) {
115                         case InetAddressUtils.IPV4:
116                             targetHost = null;
117                             targetAddress = new byte[4];
118                             for (int i = 0; i < targetAddress.length; i++) {
119                                 targetAddress[i] = input.readByte();
120                             }
121                             break;
122                         case InetAddressUtils.IPV6:
123                             targetHost = null;
124                             targetAddress = new byte[16];
125                             for (int i = 0; i < targetAddress.length; i++) {
126                                 targetAddress[i] = input.readByte();
127                             }
128                             break;
129                         case ATYP_DOMAINNAME:
130                             final int length = input.readUnsignedByte();
131                             final StringBuilder domainname = new StringBuilder();
132                             for (int i = 0; i < length; i++) {
133                                 domainname.append((char) input.readUnsignedByte());
134                             }
135                             targetHost = domainname.toString();
136                             targetAddress = null;
137                             break;
138                         default:
139                             throw new IOException("Unsupported address type: " + addressType);
140                     }
141 
142                     final int targetPort = input.readUnsignedShort();
143                     final Socket target;
144                     if (targetHost != null) {
145                         target = new Socket(targetHost, targetPort);
146                     } else {
147                         target = new Socket(InetAddress.getByAddress(targetAddress), targetPort);
148                     }
149 
150                     output.writeByte(VERSION_5);
151                     output.writeByte(0); /* success */
152                     output.writeByte(0); /* reserved */
153                     final byte[] localAddress = target.getLocalAddress().getAddress();
154                     if (localAddress.length == 4) {
155                         output.writeByte(InetAddressUtils.IPV4);
156                     } else if (localAddress.length == 16) {
157                         output.writeByte(InetAddressUtils.IPV6);
158                     } else {
159                         throw new IOException("Unsupported localAddress byte length: " + localAddress.length);
160                     }
161                     output.write(localAddress);
162                     output.writeShort(target.getLocalPort());
163                     output.flush();
164 
165                     return target;
166                 }
167 
168                 private Thread pumpStream(final InputStream input, final OutputStream output) {
169                     final Thread t = new Thread(new Runnable() {
170                         @Override
171                         public void run() {
172                             final byte[] buffer = new byte[1024 * 8];
173                             try {
174                                 while (true) {
175                                     final int read = input.read(buffer);
176                                     if (read < 0) {
177                                         break;
178                                     }
179                                     output.write(buffer, 0, read);
180                                     output.flush();
181                                 }
182                             } catch (final IOException e) {
183                             } finally {
184                                 shutdown();
185                             }
186                         }
187                     });
188                     t.start();
189                     return t;
190                 }
191 
192             }).start();
193         }
194 
195         public void shutdown() {
196             try {
197                 this.socket.close();
198             } catch (final IOException e) {
199             }
200             if (this.remote != null) {
201                 try {
202                     this.remote.close();
203                 } catch (final IOException e) {
204                 }
205             }
206         }
207 
208     }
209 
210     private final int port;
211 
212     private final List<SocksProxyHandler> handlers = new ArrayList<>();
213     private ServerSocket server;
214     private Thread serverThread;
215 
216     public SocksProxy() {
217         this(0);
218     }
219 
220     public SocksProxy(final int port) {
221         this.port = port;
222     }
223 
224     public synchronized void start() throws IOException {
225         if (this.server == null) {
226             this.server = new ServerSocket(this.port);
227             this.serverThread = new Thread(new Runnable() {
228                 @Override
229                 public void run() {
230                     try {
231                         while (true) {
232                             final Socket socket = server.accept();
233                             startSocksProxyHandler(socket);
234                         }
235                     } catch (final IOException e) {
236                     } finally {
237                         if (server != null) {
238                             try {
239                                 server.close();
240                             } catch (final IOException e) {
241                             }
242                             server = null;
243                         }
244                     }
245                 }
246             });
247             this.serverThread.start();
248         }
249     }
250 
251     public void shutdown(final TimeValue timeout) throws InterruptedException {
252         final long waitUntil = System.currentTimeMillis() + timeout.toMilliseconds();
253         Thread t = null;
254         synchronized (this) {
255             if (this.server != null) {
256                 try {
257                     this.server.close();
258                 } catch (final IOException e) {
259                 } finally {
260                     this.server = null;
261                 }
262                 t = this.serverThread;
263                 this.serverThread = null;
264             }
265             for (final SocksProxyHandler handler : this.handlers) {
266                 handler.shutdown();
267             }
268             while (!this.handlers.isEmpty()) {
269                 final long waitTime = waitUntil - System.currentTimeMillis();
270                 if (waitTime > 0) {
271                     wait(waitTime);
272                 }
273             }
274         }
275         if (t != null) {
276             final long waitTime = waitUntil - System.currentTimeMillis();
277             if (waitTime > 0) {
278                 t.join(waitTime);
279             }
280         }
281     }
282 
283     protected void startSocksProxyHandler(final Socket socket) {
284         final SocksProxyHandler handler = new SocksProxyHandler(this, socket);
285         synchronized (this) {
286             this.handlers.add(handler);
287         }
288         handler.start();
289     }
290 
291     protected synchronized void cleanupSocksProxyHandler(final SocksProxyHandler handler) {
292         this.handlers.remove(handler);
293     }
294 
295     public SocketAddress getProxyAddress() {
296         return this.server.getLocalSocketAddress();
297     }
298 
299 }