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.util.TimeValue;
42  
43  /**
44   * Cheap and nasty SOCKS protocol version 5 proxy, recommended for use in unit tests only so we can test our SOCKS client code.
45   */
46  public class SocksProxy {
47  
48      private static class SocksProxyHandler {
49  
50          public static final int VERSION_5 = 5;
51          public static final int COMMAND_CONNECT = 1;
52          public static final int ATYP_IPV4 = 1;
53          public static final int ATYP_DOMAINNAME = 3;
54          public static final int ATYP_IPV6 = 4;
55  
56          private final SocksProxy parent;
57          private final Socket socket;
58          private volatile Socket remote;
59  
60          public SocksProxyHandler(final SocksProxy parent, final Socket socket) {
61              this.parent = parent;
62              this.socket = socket;
63          }
64  
65          public void start() {
66              new Thread(new Runnable() {
67                  @Override
68                  public void run() {
69                      try {
70                          final DataInputStream input = new DataInputStream(socket.getInputStream());
71                          final DataOutputStream output = new DataOutputStream(socket.getOutputStream());
72                          final Socket target = establishConnection(input, output);
73                          remote = target;
74  
75                          final Thread t1 = pumpStream(input, target.getOutputStream());
76                          final Thread t2 = pumpStream(target.getInputStream(), output);
77                          try {
78                              t1.join();
79                          } catch (final InterruptedException e) {
80                          }
81                          try {
82                              t2.join();
83                          } catch (final InterruptedException e) {
84                          }
85                      } catch (final IOException e) {
86                      } finally {
87                          parent.cleanupSocksProxyHandler(SocksProxyHandler.this);
88                      }
89                  }
90  
91                  private Socket establishConnection(final DataInputStream input, final DataOutputStream output) throws IOException {
92                      final int clientVersion = input.readUnsignedByte();
93                      if (clientVersion != VERSION_5) {
94                          throw new IOException("SOCKS implementation only supports version 5");
95                      }
96                      final int nMethods = input.readUnsignedByte();
97                      for (int i = 0; i < nMethods; i++) {
98                          input.readUnsignedByte(); // auth method
99                      }
100                     // response
101                     output.writeByte(VERSION_5);
102                     output.writeByte(0); // no auth method
103                     output.flush();
104 
105                     input.readUnsignedByte(); // client version again
106                     final int command = input.readUnsignedByte();
107                     if (command != COMMAND_CONNECT) {
108                         throw new IOException("SOCKS implementation only supports CONNECT command");
109                     }
110                     input.readUnsignedByte(); // reserved
111 
112                     final String targetHost;
113                     final byte[] targetAddress;
114                     final int addressType = input.readUnsignedByte();
115                     switch (addressType) {
116                         case ATYP_IPV4:
117                             targetHost = null;
118                             targetAddress = new byte[4];
119                             for (int i = 0; i < targetAddress.length; i++) {
120                                 targetAddress[i] = input.readByte();
121                             }
122                             break;
123                         case ATYP_IPV6:
124                             targetHost = null;
125                             targetAddress = new byte[16];
126                             for (int i = 0; i < targetAddress.length; i++) {
127                                 targetAddress[i] = input.readByte();
128                             }
129                             break;
130                         case ATYP_DOMAINNAME:
131                             final int length = input.readUnsignedByte();
132                             final StringBuffer domainname = new StringBuffer();
133                             for (int i = 0; i < length; i++) {
134                                 domainname.append((char) input.readUnsignedByte());
135                             }
136                             targetHost = domainname.toString();
137                             targetAddress = null;
138                             break;
139                         default:
140                             throw new IOException("Unsupported address type: " + addressType);
141                     }
142 
143                     final int targetPort = input.readUnsignedShort();
144                     final Socket target;
145                     if (targetHost != null) {
146                         target = new Socket(targetHost, targetPort);
147                     } else {
148                         target = new Socket(InetAddress.getByAddress(targetAddress), targetPort);
149                     }
150 
151                     output.writeByte(VERSION_5);
152                     output.writeByte(0); /* success */
153                     output.writeByte(0); /* reserved */
154                     final byte[] localAddress = target.getLocalAddress().getAddress();
155                     if (localAddress.length == 4) {
156                         output.writeByte(ATYP_IPV4);
157                     } else if (localAddress.length == 16) {
158                         output.writeByte(ATYP_IPV6);
159                     } else {
160                         throw new IOException("Unsupported localAddress byte length: " + localAddress.length);
161                     }
162                     output.write(localAddress);
163                     output.writeShort(target.getLocalPort());
164                     output.flush();
165 
166                     return target;
167                 }
168 
169                 private Thread pumpStream(final InputStream input, final OutputStream output) {
170                     final Thread t = new Thread(new Runnable() {
171                         @Override
172                         public void run() {
173                             final byte[] buffer = new byte[1024 * 8];
174                             try {
175                                 while (true) {
176                                     final int read = input.read(buffer);
177                                     if (read < 0) {
178                                         break;
179                                     }
180                                     output.write(buffer, 0, read);
181                                     output.flush();
182                                 }
183                             } catch (final IOException e) {
184                             } finally {
185                                 shutdown();
186                             }
187                         }
188                     });
189                     t.start();
190                     return t;
191                 }
192 
193             }).start();
194         }
195 
196         public void shutdown() {
197             try {
198                 this.socket.close();
199             } catch (final IOException e) {
200             }
201             if (this.remote != null) {
202                 try {
203                     this.remote.close();
204                 } catch (final IOException e) {
205                 }
206             }
207         }
208 
209     }
210 
211     private final int port;
212 
213     private final List<SocksProxyHandler> handlers = new ArrayList<>();
214     private ServerSocket server;
215     private Thread serverThread;
216 
217     public SocksProxy() {
218         this(0);
219     }
220 
221     public SocksProxy(final int port) {
222         this.port = port;
223     }
224 
225     public synchronized void start() throws IOException {
226         if (this.server == null) {
227             this.server = new ServerSocket(this.port);
228             this.serverThread = new Thread(new Runnable() {
229                 @Override
230                 public void run() {
231                     try {
232                         while (true) {
233                             final Socket socket = server.accept();
234                             startSocksProxyHandler(socket);
235                         }
236                     } catch (final IOException e) {
237                     } finally {
238                         if (server != null) {
239                             try {
240                                 server.close();
241                             } catch (final IOException e) {
242                             }
243                             server = null;
244                         }
245                     }
246                 }
247             });
248             this.serverThread.start();
249         }
250     }
251 
252     public void shutdown(final TimeValue timeout) throws InterruptedException {
253         final long waitUntil = System.currentTimeMillis() + timeout.toMilliseconds();
254         Thread t = null;
255         synchronized (this) {
256             if (this.server != null) {
257                 try {
258                     this.server.close();
259                 } catch (final IOException e) {
260                 } finally {
261                     this.server = null;
262                 }
263                 t = this.serverThread;
264                 this.serverThread = null;
265             }
266             for (final SocksProxyHandler handler : this.handlers) {
267                 handler.shutdown();
268             }
269             while (!this.handlers.isEmpty()) {
270                 final long waitTime = waitUntil - System.currentTimeMillis();
271                 if (waitTime > 0) {
272                     wait(waitTime);
273                 }
274             }
275         }
276         if (t != null) {
277             final long waitTime = waitUntil - System.currentTimeMillis();
278             if (waitTime > 0) {
279                 t.join(waitTime);
280             }
281         }
282     }
283 
284     protected void startSocksProxyHandler(final Socket socket) {
285         final SocksProxyHandler handler = new SocksProxyHandler(this, socket);
286         synchronized (this) {
287             this.handlers.add(handler);
288         }
289         handler.start();
290     }
291 
292     protected synchronized void cleanupSocksProxyHandler(final SocksProxyHandler handler) {
293         this.handlers.remove(handler);
294     }
295 
296     public SocketAddress getProxyAddress() {
297         return this.server.getLocalSocketAddress();
298     }
299 
300 }