1 /*

     2  * Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved.

     3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

     4  *

     5  * This code is free software; you can redistribute it and/or modify it

     6  * under the terms of the GNU General Public License version 2 only, as

     7  * published by the Free Software Foundation.

     8  *

     9  * This code is distributed in the hope that it will be useful, but WITHOUT

    10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

    11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

    12  * version 2 for more details (a copy is included in the LICENSE file that

    13  * accompanied this code).

    14  *

    15  * You should have received a copy of the GNU General Public License version

    16  * 2 along with this work; if not, write to the Free Software Foundation,

    17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

    18  *

    19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

    20  * or visit www.oracle.com if you need additional information or have any

    21  * questions.

    22  */

    23 

    24 import java.io.ByteArrayOutputStream;

    25 import java.io.Closeable;

    26 import java.io.IOException;

    27 import java.io.InputStream;

    28 import java.io.OutputStream;

    29 import java.net.InetAddress;

    30 import java.net.ServerSocket;

    31 import java.net.Socket;

    32 import java.util.ArrayList;

    33 import java.util.Arrays;

    34 import java.util.HashSet;

    35 import java.util.List;

    36 import java.util.Objects;

    37 import java.util.Set;

    38 import java.util.concurrent.ExecutorService;

    39 import java.util.concurrent.Executors;

    40 import java.util.concurrent.RejectedExecutionException;

    41 import java.util.concurrent.ThreadFactory;

    42 import java.util.concurrent.TimeUnit;

    43 import java.util.stream.Stream;

    44 

    45 import static java.lang.StackWalker.Option.RETAIN_CLASS_REFERENCE;

    46 

    47 /**

    48  * A base dummy ldap server.

    49  *

    50  * For any ldap tests which required a simple dummy server to support the

    51  * test, may use this server directly with specifying ConnectionHandler,

    52  * SessionHandler/RequestHandler, or extends to build more complex server logic.

    53  *

    54  * This server already extends Thread and implements AutoCloseable, so it can

    55  * be started in thread and integrated with try-with-resources

    56  *

    57  * To initiate a instance of this server, valid ServerSocket could be supplied,

    58  * it will allow the flexibility for listening address/port customization

    59  * and SSL usage, for default no parameter constructor, a ServerSocket which

    60  * listen on loopback address will be created.

    61  *

    62  * To use this dummy server in test, user could customize the processing logic

    63  * in three level with below handler interface.

    64  *   -ConnectionHandler provide connection level handling, server will hand

    65  *                      over accepted socket and processing thread to handler.

    66  *                      By default, DefaultConnectionHandler will be used if no

    67  *                      specified, it reads full ldap request message then

    68  *                      pass it to RequestHandler instance which returned by

    69  *                      SessionHandler per session.

    70  *

    71  *   -SessionHandler    provide session level handling when DefaultConnectionHandler

    72  *                      been used, it's to retrieve RequestHandler instance of

    73  *                      current session.

    74  *                      For most of tests, only one session need to be handled

    75  *                      on server or all ldap request could be handled by same

    76  *                      logic whatever current session is, user can use

    77  *                      setCommonRequestHandler to setup one single session

    78  *                      handler which will always return given RequestHandler

    79  *                      instance.

    80  *

    81  *   -RequestHandler    provide ldap message request handling when

    82  *                      DefaultConnectionHandler been used.

    83  *

    84  * @see ConnectionHandler

    85  * @see SessionHandler

    86  * @see RequestHandler

    87  */

    88 public class BaseLdapServer extends Thread implements AutoCloseable {

    89     private volatile boolean isRunning;

    90     private final List<Socket> socketList = new ArrayList<>();

    91     private ServerSocket serverSocket;

    92     private ExecutorService workingPool;

    93     private ConnectionHandler connectionHandler;

    94     private SessionHandler sessionHandler;

    95     private boolean useDaemonThread = false;

    96 

    97     enum DebugLevel {

    98         FULL,      // all debug message will be printed

    99         NONE,      // none of debug message will be printed

   100         CUSTOMIZE  // only specified class debug message will be printed

   101     }

   102 

   103     private StackWalker stackWalker = null;

   104     private DebugLevel debugLevel = DebugLevel.NONE;

   105     private Set<Class<?>> debugOptions = new HashSet<>();

   106 

   107     /**

   108      * BaseLdapServer overload default constructor.

   109      *

   110      * @throws IOException if an I/O error occurs when opening the socket.

   111      */

   112     public BaseLdapServer() throws IOException {

   113         this(new ServerSocket(0, 0, InetAddress.getLoopbackAddress()));

   114     }

   115 

   116     /**

   117      * BaseLdapServer overload constructor with given server socket.

   118      *

   119      * @param serverSocket given server socket

   120      */

   121     public BaseLdapServer(ServerSocket serverSocket) {

   122         this(serverSocket, false);

   123     }

   124 

   125     /**

   126      * BaseLdapServer constructor with given server socket and specify whether

   127      * use daemon for each accept connection handling thread.

   128      *

   129      * @param serverSocket    given server socket

   130      * @param useDaemonThread <tt>true</tt> if use daemon thread

   131      */

   132     public BaseLdapServer(ServerSocket serverSocket, boolean useDaemonThread) {

   133         this.serverSocket = Objects.requireNonNull(serverSocket);

   134         this.useDaemonThread = useDaemonThread;

   135         if (useDaemonThread) {

   136             workingPool = Executors

   137                     .newCachedThreadPool(new DefaultDaemonThreadFactory());

   138         } else {

   139             workingPool = Executors.newCachedThreadPool();

   140         }

   141         try {

   142             stackWalker = StackWalker.getInstance(RETAIN_CLASS_REFERENCE);

   143         } catch (SecurityException se) {

   144             // just ignore

   145         }

   146     }

   147 

   148     @Override

   149     public void run() {

   150         if (getConnectionHandler() == null) {

   151             debug("INFO: No connection handler been specified, try default.");

   152             connectionHandler = new DefaultConnectionHandler();

   153         }

   154         debug("INFO: Using connection handler : " + getConnectionHandler()

   155                 .getClass().getName());

   156         debug("INFO: LdapServer running and listening on port " + getPort());

   157         try {

   158             while (isRunning) {

   159                 Socket socket = serverSocket.accept();

   160                 debug("INFO: Accept new connection " + socket);

   161                 synchronized (socketList) {

   162                     socketList.add(socket);

   163                 }

   164                 workingPool.submit(() -> getConnectionHandler()

   165                         .handleConnection(socket));

   166             }

   167         } catch (IOException | RejectedExecutionException e) {

   168             if (isRunning) {

   169                 throw new RuntimeException(e);

   170             } else {

   171                 debug("INFO: Server exit.");

   172             }

   173         }

   174     }

   175 

   176     /*

   177      * Override Thread.start()

   178      */

   179     @Override

   180     public synchronized void start() {

   181         super.start();

   182         isRunning = true;

   183     }

   184 

   185     /**

   186      * Start Server thread and return itself for method chaining

   187      *

   188      * @return current server instance

   189      */

   190     @SuppressWarnings("unchecked")

   191     public <T extends BaseLdapServer> T startServer() {

   192         start();

   193         return (T) this;

   194     }

   195 

   196     /**

   197      * Stop server.

   198      */

   199     public void stopServer() {

   200         debug("INFO: Stopping Server.");

   201         isRunning = false;

   202         workingPool.shutdown();

   203         cleanupClosableRes(serverSocket);

   204         if (!useDaemonThread) {

   205             // let's cleanup thread pool

   206             synchronized (socketList) {

   207                 socketList.forEach(BaseLdapServer::cleanupClosableRes);

   208             }

   209             try {

   210                 if (!workingPool.awaitTermination(10, TimeUnit.SECONDS)) {

   211                     workingPool.shutdownNow();

   212                 }

   213             } catch (InterruptedException e) {

   214                 workingPool.shutdownNow();

   215                 Thread.currentThread().interrupt();

   216             }

   217         }

   218     }

   219 

   220     /**

   221      * Return local port which server is listening.

   222      *

   223      * @return port which server is listening

   224      */

   225     public int getPort() {

   226         if (serverSocket != null) {

   227             return serverSocket.getLocalPort();

   228         } else {

   229             return -1;

   230         }

   231     }

   232 

   233     /**

   234      * Return flag to indicate whether current server is running.

   235      *

   236      * @return <tt>true</tt> if current server is running.

   237      */

   238     public boolean isRunning() {

   239         return isRunning;

   240     }

   241 

   242     /**

   243      * Return ConnectionHandler instance

   244      *

   245      * @return ConnectionHandler instance

   246      * @see ConnectionHandler

   247      */

   248     ConnectionHandler getConnectionHandler() {

   249         return connectionHandler;

   250     }

   251 

   252     /**

   253      * Set ConnectionHandler when server is not running.

   254      *

   255      * @param connHandler ConnectionHandler instance

   256      * @return current server instance for method chaining

   257      */

   258     @SuppressWarnings("unchecked")

   259     public <T extends BaseLdapServer> T setConnectionHandler(

   260             ConnectionHandler connHandler) {

   261         if (!isRunning) {

   262             connectionHandler = connHandler;

   263         }

   264 

   265         return (T) this;

   266     }

   267 

   268     /**

   269      * Return SessionHandler instance

   270      *

   271      * @return SessionHandler instance

   272      * @see SessionHandler

   273      */

   274     SessionHandler getSessionHandler() {

   275         return sessionHandler;

   276     }

   277 

   278     /**

   279      * Set SessionHandler when server is not running.

   280      *

   281      * @param sessionHandler given SessionHandler

   282      * @return current server instance for method chaining

   283      */

   284     @SuppressWarnings("unchecked")

   285     public <T extends BaseLdapServer> T setSessionHandler(

   286             SessionHandler sessionHandler) {

   287         if (!isRunning) {

   288             this.sessionHandler = sessionHandler;

   289         }

   290 

   291         return (T) this;

   292     }

   293 

   294     /**

   295      * Set one common RequestHandler, it will be used to handle all requests

   296      * whatever current session is.

   297      *

   298      * For most of tests, server only need to handle one session, use this

   299      * method will create stateless session handler with given request handler.

   300      *

   301      * @param requestHandler RequestHandler instance

   302      * @return current server instance for method chaining

   303      */

   304     @SuppressWarnings("unchecked")

   305     public <T extends BaseLdapServer> T setCommonRequestHandler(

   306             RequestHandler requestHandler) {

   307         if (!isRunning) {

   308             // ignore any session, always return fixed request handler

   309             setSessionHandler(socket -> requestHandler);

   310         }

   311 

   312         return (T) this;

   313     }

   314 

   315     @Override

   316     public void close() {

   317         stopServer();

   318     }

   319 

   320     /**

   321      * Cleanup any given closable resource

   322      *

   323      * @param res given closable resource

   324      */

   325     static void cleanupClosableRes(Closeable res) {

   326         if (res != null) {

   327             try {

   328                 res.close();

   329             } catch (IOException e) {

   330                 // ignore

   331             }

   332         }

   333     }

   334 

   335     /**

   336      * Set debug level to specify which kinds of debug message will be printed.

   337      *

   338      * @param debugLevel given debug level

   339      * @param opts       given opts if debug level is DebugLevel.CUSTOMIZE

   340      * @return current server instance for method chaining

   341      */

   342     @SuppressWarnings("unchecked")

   343     public <T extends BaseLdapServer> T setDebugLevel(DebugLevel debugLevel,

   344             Class<?>... opts) {

   345         Objects.requireNonNull(debugLevel);

   346         if (!isRunning) {

   347             this.debugLevel = debugLevel;

   348             if (debugLevel == DebugLevel.CUSTOMIZE) {

   349                 debugOptions.clear();

   350                 Stream.of(opts).filter(Objects::nonNull)

   351                         .forEach(debugOptions::add);

   352             }

   353         }

   354 

   355         return (T) this;

   356     }

   357 

   358     /**

   359      * Print given message if debug enabled.

   360      *

   361      * @param message given message to print

   362      */

   363     void debug(String message) {

   364         switch (debugLevel) {

   365             case FULL:

   366                 System.out.println((stackWalker != null ?

   367                         stackWalker.getCallerClass().getName() :

   368                         "") + ": " + message);

   369                 break;

   370             case CUSTOMIZE:

   371                 if (stackWalker != null) {

   372                     if (debugOptions.contains(stackWalker.getCallerClass())) {

   373                         System.out.println(

   374                                 stackWalker.getCallerClass().getName() + ": "

   375                                         + message);

   376                     }

   377                 }

   378                 break;

   379             case NONE:

   380             default:

   381                 break;

   382         }

   383     }

   384 

   385     class DefaultDaemonThreadFactory implements ThreadFactory {

   386 

   387         private ThreadFactory defaultThreadFactory = Executors

   388                 .defaultThreadFactory();

   389 

   390         @Override

   391         public Thread newThread(Runnable r) {

   392             Thread thread = defaultThreadFactory.newThread(r);

   393             thread.setDaemon(true);

   394             return thread;

   395         }

   396     }

   397 

   398     /**

   399      * Default connection handler implementation.

   400      */

   401     class DefaultConnectionHandler implements ConnectionHandler {

   402         @Override

   403         public void handleConnection(Socket socket) {

   404             try (socket;

   405                     OutputStream out = socket.getOutputStream();

   406                     InputStream in = socket.getInputStream()) {

   407                 byte[] inBuffer = new byte[1024];

   408                 int count;

   409                 byte[] request;

   410 

   411                 ByteArrayOutputStream buffer = new ByteArrayOutputStream();

   412                 int msgLen = -1;

   413 

   414                 while ((count = in.read(inBuffer)) > 0) {

   415                     buffer.write(inBuffer, 0, count);

   416                     if (msgLen <= 0) {

   417                         msgLen = getMessageLength(buffer.toByteArray());

   418                     }

   419 

   420                     if (msgLen > 0 && buffer.size() >= msgLen) {

   421                         if (buffer.size() > msgLen) {

   422                             byte[] tmpBuffer = buffer.toByteArray();

   423                             request = Arrays.copyOf(tmpBuffer, msgLen);

   424                             buffer.reset();

   425                             buffer.write(tmpBuffer, msgLen,

   426                                     tmpBuffer.length - msgLen);

   427                         } else {

   428                             request = buffer.toByteArray();

   429                             buffer.reset();

   430                         }

   431                         msgLen = -1;

   432                     } else {

   433                         debug("INFO: request msg not complete, received "

   434                                 + buffer.size() + ", expected " + msgLen);

   435                         continue;

   436                     }

   437 

   438                     if (getSessionHandler() != null) {

   439                         var handler = getSessionHandler()

   440                                 .getRequestHandler(socket);

   441                         if (handler != null) {

   442                             debug("INFO: Process request. Session handler : "

   443                                     + getSessionHandler()

   444                                     + ", Request handler : " + handler);

   445                             handler.handleRequest(new LdapMessage(request),

   446                                     out);

   447                         } else {

   448                             debug("WARNING: no valid request handler returned from "

   449                                     + getSessionHandler() + ", " + socket);

   450                         }

   451                     } else {

   452                         debug("WARNING: no valid session handler been specified, discard request.");

   453                     }

   454                 }

   455                 debug("INFO: Connection Handler exit.");

   456             } catch (IOException e) {

   457                 if (!isRunning()) {

   458                     debug("INFO: Connection Handler exit : " + e.getMessage());

   459                 } else {

   460                     e.printStackTrace();

   461                 }

   462             }

   463         }

   464 

   465         private int getMessageLength(byte[] encoding) {

   466             if (encoding.length < 2) {

   467                 // no enough data to extract msg len, just return -1

   468                 return -1;

   469             }

   470 

   471             if (encoding[0] != 0x30) {

   472                 throw new RuntimeException("Error: bad LDAP encoding message: "

   473                         + "expected ASN.1 SEQUENCE tag (0x30), encountered "

   474                         + encoding[0]);

   475             }

   476 

   477             int len;

   478             int index = 1;

   479             int payloadLen = 0;

   480 

   481             if ((encoding[1] & 0x80) == 0x80) {

   482                 len = (encoding[1] & 0x0F);

   483                 index++;

   484             } else {

   485                 len = 1;

   486             }

   487 

   488             if (len > 4) {

   489                 throw new RuntimeException(

   490                         "Error: LDAP encoding message payload too large");

   491             }

   492 

   493             if (encoding.length < index + len) {

   494                 // additional data required to extract payload len, return -1

   495                 return -1;

   496             }

   497 

   498             for (byte b : Arrays.copyOfRange(encoding, index, index + len)) {

   499                 payloadLen = payloadLen << 8 | (b & 0xFF);

   500             }

   501 

   502             if (payloadLen <= 0) {

   503                 throw new RuntimeException(

   504                         "Error: invalid LDAP encoding message length or payload too large");

   505             }

   506 

   507             return index + len + payloadLen;

   508         }

   509     }

   510 }