/*
    * 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.eclipse.aether.named.ipc;
  
  import java.io.BufferedReader;
  import java.io.DataInputStream;
  import java.io.DataOutputStream;
  import java.io.IOException;
  import java.io.InputStreamReader;
  import java.net.SocketAddress;
  import java.nio.channels.ByteChannel;
  import java.nio.channels.Channels;
  import java.nio.channels.ServerSocketChannel;
  import java.nio.channels.SocketChannel;
  import java.nio.charset.StandardCharsets;
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.CopyOnWriteArrayList;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  
  /**
   * Implementation of the server side.
   * The server instance is bound to a given maven repository.
   *
   * @since 2.0.1
   */
  public class IpcServer {
      /**
       * Should the IPC server not fork? (i.e. for testing purposes)
       *
       * @configurationSource {@link System#getProperty(String, String)}
       * @configurationType {@link java.lang.Boolean}
       * @configurationDefaultValue {@link #DEFAULT_NO_FORK}
       */
      public static final String SYSTEM_PROP_NO_FORK = "aether.named.ipc.nofork";
  
      public static final boolean DEFAULT_NO_FORK = false;
  
      /**
       * IPC idle timeout in seconds. If there is no IPC request during idle time, it will stop.
       *
       * @configurationSource {@link System#getProperty(String, String)}
       * @configurationType {@link java.lang.Integer}
       * @configurationDefaultValue {@link #DEFAULT_IDLE_TIMEOUT}
       */
      public static final String SYSTEM_PROP_IDLE_TIMEOUT = "aether.named.ipc.idleTimeout";
  
      public static final int DEFAULT_IDLE_TIMEOUT = 300;
  
      /**
       * IPC socket family to use.
       *
       * @configurationSource {@link System#getProperty(String, String)}
       * @configurationType {@link java.lang.String}
       * @configurationDefaultValue {@link #DEFAULT_FAMILY}
       */
      public static final String SYSTEM_PROP_FAMILY = "aether.named.ipc.family";
  
      public static final String DEFAULT_FAMILY = "unix";
  
      /**
       * Should the IPC server not use native executable?
       *
       * @configurationSource {@link System#getProperty(String, String)}
       * @configurationType {@link java.lang.Boolean}
       * @configurationDefaultValue {@link #DEFAULT_NO_NATIVE}
       */
      public static final String SYSTEM_PROP_NO_NATIVE = "aether.named.ipc.nonative";
  
      public static final boolean DEFAULT_NO_NATIVE = true;
  
      /**
       * The name if the IPC server native executable (without file extension like ".exe")
       *
       * @configurationSource {@link System#getProperty(String, String)}
       * @configurationType {@link java.lang.String}
       * @configurationDefaultValue {@link #DEFAULT_NATIVE_NAME}
       */
     public static final String SYSTEM_PROP_NATIVE_NAME = "aether.named.ipc.nativeName";
 
     public static final String DEFAULT_NATIVE_NAME = "ipc-sync";
 
     /**
      * Should the IPC server log debug messages? (i.e. for testing purposes)
      *
      * @configurationSource {@link System#getProperty(String, String)}
      * @configurationType {@link java.lang.Boolean}
      * @configurationDefaultValue {@link #DEFAULT_DEBUG}
      */
     public static final String SYSTEM_PROP_DEBUG = "aether.named.ipc.debug";
 
     public static final boolean DEFAULT_DEBUG = false;
 
     private final ServerSocketChannel serverSocket;
     private final Map<SocketChannel, Thread> clients = new ConcurrentHashMap<>();
     private final AtomicInteger counter = new AtomicInteger();
     private final Map<String, Lock> locks = new ConcurrentHashMap<>();
     private final Map<String, Context> contexts = new ConcurrentHashMap<>();
     private static final boolean DEBUG =
             Boolean.parseBoolean(System.getProperty(SYSTEM_PROP_DEBUG, Boolean.toString(DEFAULT_DEBUG)));
     private final long idleTimeout;
     private final String bootstrapToken;
     private volatile long lastUsed;
     private volatile boolean closing;
 
     /**
      * @deprecated a server created without a bootstrap token refuses {@link IpcMessages#REQUEST_STOP} requests;
      * use {@link #IpcServer(SocketFamily, String)} instead.
      */
     @Deprecated
     public IpcServer(SocketFamily family) throws IOException {
         this(family, null);
     }
 
     /**
      * Creates a server that honors a remote stop request only when it carries the given bootstrap token, which is
      * shared exclusively with the client that spawned this server. The rest of the protocol is unauthenticated,
      * but destructive cross-client operations (closing foreign contexts, stopping the daemon) are refused.
      *
      * @since 2.0.23
      */
     public IpcServer(SocketFamily family, String bootstrapToken) throws IOException {
         this.bootstrapToken = bootstrapToken;
         serverSocket = family.openServerSocket();
         long timeout = TimeUnit.SECONDS.toNanos(DEFAULT_IDLE_TIMEOUT);
         String str = System.getProperty(SYSTEM_PROP_IDLE_TIMEOUT);
         if (str != null) {
             try {
                 TimeUnit unit = TimeUnit.SECONDS;
                 if (str.endsWith("ms")) {
                     unit = TimeUnit.MILLISECONDS;
                     str = str.substring(0, str.length() - 2);
                 }
                 long dur = Long.parseLong(str);
                 timeout = unit.toNanos(dur);
             } catch (NumberFormatException e) {
                 error("Property " + SYSTEM_PROP_IDLE_TIMEOUT + " specified with invalid value: " + str, e);
             }
         }
         idleTimeout = timeout;
     }
 
     public static void main(String[] args) throws Exception {
         // When spawning a new process, the child process is create within
         // the same process group.  This means that a few signals are sent
         // to the whole group.  This is the case for SIGINT (Ctrl-C) and
         // SIGTSTP (Ctrl-Z) which are both sent to all the processed in the
         // group when initiated from the controlling terminal.
         // This is only a problem when the client creates the daemon, but
         // without ignoring those signals, a client being interrupted will
         // also interrupt and kill the daemon.
         try {
             sun.misc.Signal.handle(new sun.misc.Signal("INT"), sun.misc.SignalHandler.SIG_IGN);
             if (!IpcClient.IS_WINDOWS) {
                 sun.misc.Signal.handle(new sun.misc.Signal("TSTP"), sun.misc.SignalHandler.SIG_IGN);
             }
         } catch (Throwable t) {
             error("Unable to ignore INT and TSTP signals", t);
         }
 
         String family = args[0];
         String tmpAddress = args[1];
         String rand = args[2];
         if ("-".equals(rand)) {
             // the bootstrap token is passed via stdin instead of argv: process arguments are commonly visible
             // to other local users (e.g. /proc/<pid>/cmdline), and this token authorizes stopping the daemon
             BufferedReader reader = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
             rand = reader.readLine();
             if (rand == null || rand.isEmpty()) {
                 throw new IOException("Expected the bootstrap token on standard input");
             }
         }
 
         runServer(SocketFamily.valueOf(family), tmpAddress, rand);
     }
 
     static IpcServer runServer(SocketFamily family, String tmpAddress, String rand) throws IOException {
         IpcServer server = new IpcServer(family, rand);
         run(server::run, false); // this is one-off
         String address = SocketFamily.toString(server.getLocalAddress());
         SocketAddress socketAddress = SocketFamily.fromString(tmpAddress);
         try (SocketChannel socket = SocketChannel.open(socketAddress)) {
             try (DataOutputStream dos = new DataOutputStream(Channels.newOutputStream(socket))) {
                 dos.writeUTF(rand);
                 dos.writeUTF(address);
                 dos.flush();
             }
         }
 
         return server;
     }
 
     private static void debug(String msg, Object... args) {
         if (DEBUG) {
             System.out.printf("[ipc] [debug] " + msg + "\n", args);
         }
     }
 
     private static void info(String msg, Object... args) {
         System.out.printf("[ipc] [info] " + msg + "\n", args);
     }
 
     private static void error(String msg, Throwable t) {
         System.out.println("[ipc] [error] " + msg);
         t.printStackTrace(System.out);
     }
 
     private static void run(Runnable runnable, boolean daemon) {
         Thread thread = new Thread(runnable);
         if (daemon) {
             thread.setDaemon(true);
         }
         thread.start();
     }
 
     public SocketAddress getLocalAddress() throws IOException {
         return serverSocket.getLocalAddress();
     }
 
     public void run() {
         try {
             info("IpcServer started at %s", getLocalAddress().toString());
             use();
             run(this::expirationCheck, true);
             while (!closing) {
                 SocketChannel socket = this.serverSocket.accept();
                 run(() -> client(socket), false);
             }
         } catch (Throwable t) {
             if (!closing) {
                 error("Error running sync server loop", t);
             }
         }
     }
 
     private void client(SocketChannel socket) {
         int c;
         synchronized (clients) {
             clients.put(socket, Thread.currentThread());
             c = clients.size();
         }
         info("New client connected (%d connected)", c);
         use();
         Map<String, Context> clientContexts = new ConcurrentHashMap<>();
         try {
             ByteChannel wrapper = new ByteChannelWrapper(socket);
             DataInputStream input = new DataInputStream(Channels.newInputStream(wrapper));
             DataOutputStream output = new DataOutputStream(Channels.newOutputStream(wrapper));
             while (!closing) {
                 int requestId = input.readInt();
                 int sz = input.readInt();
                 if (sz < 0) {
                     throw new IOException("Received invalid request size: " + sz);
                 }
                 // do not preallocate from an unauthenticated wire-supplied size; grow with actually received data
                 List<String> request = new ArrayList<>(Math.min(sz, 1024));
                 for (int i = 0; i < sz; i++) {
                     request.add(input.readUTF());
                 }
                 if (request.isEmpty()) {
                     throw new IOException("Received invalid request");
                 }
                 use();
                 String contextId;
                 Context context;
                 String command = request.remove(0);
                 switch (command) {
                     case IpcMessages.REQUEST_CONTEXT:
                         if (request.size() != 1) {
                             throw new IOException("Expected one argument for " + command + " but got " + request);
                         }
                         boolean shared = Boolean.parseBoolean(request.remove(0));
                         context = new Context(shared);
                         contexts.put(context.id, context);
                         clientContexts.put(context.id, context);
                         synchronized (output) {
                             debug("Created context %s", context.id);
                             output.writeInt(requestId);
                             output.writeInt(2);
                             output.writeUTF(IpcMessages.RESPONSE_CONTEXT);
                             output.writeUTF(context.id);
                             output.flush();
                         }
                         break;
                     case IpcMessages.REQUEST_ACQUIRE:
                         if (request.isEmpty()) {
                             throw new IOException(
                                     "Expected at least one argument for " + command + " but got " + request);
                         }
                         contextId = request.remove(0);
                         // contexts are scoped per connection: a client may only use contexts it created itself
                         context = clientContexts.get(contextId);
                         if (context == null) {
                             throw new IOException(
                                     "Unknown context: " + contextId + ". Known contexts = " + clientContexts.keySet());
                         }
                         context.lock(request).thenRun(() -> sendAcquireResponse(output, socket, requestId, context));
                         break;
                     case IpcMessages.REQUEST_CLOSE:
                         if (request.size() != 1) {
                             throw new IOException("Expected one argument for " + command + " but got " + request);
                         }
                         contextId = request.remove(0);
                         // contexts are scoped per connection: a client may only close contexts it created itself
                         context = clientContexts.remove(contextId);
                         if (context == null) {
                             throw new IOException(
                                     "Unknown context: " + contextId + ". Known contexts = " + clientContexts.keySet());
                         }
                         contexts.remove(contextId);
                         context.unlock();
                         synchronized (output) {
                             debug("Closing context %s", context.id);
                             output.writeInt(requestId);
                             output.writeInt(1);
                             output.writeUTF(IpcMessages.RESPONSE_CLOSE);
                             output.flush();
                         }
                         break;
                     case IpcMessages.REQUEST_STOP:
                         if (request.size() > 1) {
                             throw new IOException(
                                     "Expected at most one argument for " + command + " but got " + request);
                         }
                         String stopToken = request.isEmpty() ? null : request.remove(0);
                         if (bootstrapToken == null || !bootstrapToken.equals(stopToken)) {
                             // the protocol is otherwise unauthenticated: only the client that spawned this
                             // server (and thus knows the bootstrap token) may stop it for everybody else
                             throw new IOException("Stop request rejected: missing or invalid bootstrap token");
                         }
                         synchronized (output) {
                             debug("Stopping server");
                             output.writeInt(requestId);
                             output.writeInt(1);
                             output.writeUTF(IpcMessages.RESPONSE_STOP);
                             output.flush();
                         }
                         close();
                         break;
                     default:
                         throw new IOException("Unknown request: " + request.get(0));
                 }
             }
         } catch (Throwable t) {
             if (!closing) {
                 error("Error processing request", t);
             }
         } finally {
             if (!closing) {
                 info("Client disconnecting...");
             }
             clientContexts.values().forEach(context -> {
                 contexts.remove(context.id);
                 context.unlock();
             });
             try {
                 socket.close();
             } catch (IOException ioException) {
                 // ignore
             }
             synchronized (clients) {
                 clients.remove(socket);
                 c = clients.size();
             }
             if (!closing) {
                 info("%d clients remained", c);
             }
         }
     }
 
     private void sendAcquireResponse(DataOutputStream output, SocketChannel socket, int requestId, Context context) {
         try {
             synchronized (output) {
                 debug("Locking in context %s", context.id);
                 output.writeInt(requestId);
                 output.writeInt(1);
                 output.writeUTF(IpcMessages.RESPONSE_ACQUIRE);
                 output.flush();
             }
         } catch (IOException e) {
             try {
                 socket.close();
             } catch (IOException ioException) {
                 e.addSuppressed(ioException);
             }
             error("Error writing lock response", e);
         }
     }
 
     private void use() {
         lastUsed = System.nanoTime();
     }
 
     private void expirationCheck() {
         while (true) {
             long current = System.nanoTime();
             long left = (lastUsed + idleTimeout) - current;
             if (clients.isEmpty() && left < 0) {
                 info("IpcServer expired, closing");
                 close();
                 break;
             } else {
                 try {
                     Thread.sleep(Math.max(1, TimeUnit.NANOSECONDS.toMillis(left)));
                 } catch (InterruptedException e) {
                     info("IpcServer expiration check interrupted, closing");
                     close();
                     break;
                 }
             }
         }
     }
 
     void close() {
         closing = true;
         try {
             serverSocket.close();
         } catch (IOException e) {
             error("Error closing server socket", e);
         }
         clients.forEach((s, t) -> {
             try {
                 s.close();
             } catch (IOException e) {
                 // ignore
             }
             t.interrupt();
         });
     }
 
     static class Waiter {
         final Context context;
         final CompletableFuture<Void> future;
 
         Waiter(Context context, CompletableFuture<Void> future) {
             this.context = context;
             this.future = future;
         }
     }
 
     static class Lock {
 
         final String key;
 
         List<Context> holders;
         List<Waiter> waiters;
 
         Lock(String key) {
             this.key = key;
         }
 
         public synchronized CompletableFuture<Void> lock(Context context) {
             if (holders == null) {
                 holders = new ArrayList<>();
             }
             if (holders.isEmpty() || holders.get(0).shared && context.shared) {
                 holders.add(context);
                 return CompletableFuture.completedFuture(null);
             }
             if (waiters == null) {
                 waiters = new ArrayList<>();
             }
 
             CompletableFuture<Void> future = new CompletableFuture<>();
             waiters.add(new Waiter(context, future));
             return future;
         }
 
         public void unlock(Context context) {
             List<CompletableFuture<Void>> toComplete;
             synchronized (this) {
                 toComplete = new ArrayList<>();
                 if (holders.remove(context)) {
                     while (waiters != null
                             && !waiters.isEmpty()
                             && (holders.isEmpty() || holders.get(0).shared && waiters.get(0).context.shared)) {
                         Waiter waiter = waiters.remove(0);
                         holders.add(waiter.context);
                         toComplete.add(waiter.future);
                     }
                 } else if (waiters != null) {
                     for (Iterator<Waiter> it = waiters.iterator(); it.hasNext(); ) {
                         Waiter waiter = it.next();
                         if (waiter.context == context) {
                             it.remove();
                             waiter.future.cancel(false);
                         }
                     }
                 }
             }
             toComplete.forEach(f -> f.complete(null));
         }
 
         public synchronized boolean isEmpty() {
             return (holders == null || holders.isEmpty()) && (waiters == null || waiters.isEmpty());
         }
     }
 
     class Context {
 
         final String id;
         final boolean shared;
         final List<String> locks = new CopyOnWriteArrayList<>();
 
         Context(boolean shared) {
             this.id = String.format("%08x", counter.incrementAndGet());
             this.shared = shared;
         }
 
         public CompletableFuture<?> lock(List<String> keys) {
             locks.addAll(keys);
             CompletableFuture<?>[] futures = keys.stream()
                     .map(k -> IpcServer.this.locks.computeIfAbsent(k, Lock::new))
                     .map(l -> l.lock(this))
                     .toArray(CompletableFuture[]::new);
             return CompletableFuture.allOf(futures);
         }
 
         public void unlock() {
             locks.stream()
                     .map(k -> IpcServer.this.locks.computeIfAbsent(k, Lock::new))
                     .forEach(l -> {
                         l.unlock(this);
                         IpcServer.this.locks.compute(l.key, (k, v) -> (v == l && v.isEmpty()) ? null : v);
                     });
         }
     }
 }