/*
 * junixsocket
 *
 * Copyright 2009-2024 Christian Kohlschütter
 *
 * Licensed 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.newsclub.net.unix.vsock;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.SocketException;
import java.nio.charset.StandardCharsets;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.newsclub.net.unix.AFSocket;
import org.newsclub.net.unix.AFSocketAddress;
import org.newsclub.net.unix.AFSocketConnector;
import org.newsclub.net.unix.AFUNIXSocket;
import org.newsclub.net.unix.AFUNIXSocketAddress;
import org.newsclub.net.unix.AFVSOCKSocketAddress;
import org.newsclub.net.unix.AddressUnavailableSocketException;

/**
 * Provides access to AF_VSOCK connections that aren't directly accessible but exposed via a
 * proxying/multiplexing Unix domain socket.
 *
 * @author Christian Kohlschütter
 * @see #openFirecrackerStyleConnector(AFUNIXSocketAddress, int)
 * @see #openDirectConnector()
 */
public final class AFVSOCKProxySocketConnector implements
    AFSocketConnector<AFVSOCKSocketAddress, AFSocketAddress> {
  private static final AFSocketConnector<AFVSOCKSocketAddress, AFSocketAddress> DIRECT_CONNECTOR =
      new AFSocketConnector<AFVSOCKSocketAddress, AFSocketAddress>() {

        @Override
        public AFSocket<? extends AFSocketAddress> connect(AFVSOCKSocketAddress addr)
            throws IOException {
          return addr.newConnectedSocket();
        }
      };

  private static final Pattern PAT_OK = Pattern.compile("OK ([0-9]+)");
  private static final byte[] OK = {'O', 'K', ' '};
  private final AFUNIXSocketAddress connectorAddress;
  private final int allowedCID;

  private AFVSOCKProxySocketConnector(AFUNIXSocketAddress connectorAddress, int allowedCID) {
    this.connectorAddress = connectorAddress;
    this.allowedCID = allowedCID;
  }

  /**
   * Returns an instance that is configured to support
   * [Firecracker-style](https://github.com/firecracker-microvm/firecracker/blob/main/docs/vsock.md)
   * Unix domain sockets.
   *
   * @param connectorAddress The unix socket address pointing at the Firecracker-style multiplexing
   *          domain socket.
   * @param allowedCID The permitted CID, or {@link AFVSOCKSocketAddress#VMADDR_CID_ANY} for "any".
   * @return The instance.
   */
  public static AFSocketConnector<AFVSOCKSocketAddress, AFSocketAddress> openFirecrackerStyleConnector(
      AFUNIXSocketAddress connectorAddress, int allowedCID) {
    return new AFVSOCKProxySocketConnector(connectorAddress, allowedCID);
  }

  /**
   * Returns an instance that is configured to connect directly to the given address.
   *
   * @return The direct instance.
   */
  public static AFSocketConnector<AFVSOCKSocketAddress, AFSocketAddress> openDirectConnector() {
    return DIRECT_CONNECTOR;
  }

  /**
   * Connects to the given AF_VSOCK address.
   *
   * @param vsockAddress The address to connect to.
   * @return The connected socket.
   * @throws IOException on error.
   * @throws AddressUnavailableSocketException if the CID is not covered by this connector.
   */
  @Override
  @SuppressWarnings("Finally" /* errorprone */)
  public AFSocket<?> connect(AFVSOCKSocketAddress vsockAddress) throws IOException {
    int cid = vsockAddress.getVSOCKCID();
    if (cid != allowedCID && cid != AFVSOCKSocketAddress.VMADDR_CID_ANY
        && allowedCID != AFVSOCKSocketAddress.VMADDR_CID_ANY) {
      throw new AddressUnavailableSocketException("Connector does not cover CID " + cid);
    }

    @SuppressWarnings("resource")
    AFUNIXSocket sock = connectorAddress.newConnectedSocket();
    InputStream in = sock.getInputStream();
    OutputStream out = sock.getOutputStream();

    boolean success = false;

    try { // NOPMD.UseTryWithResources
      String connectLine = "CONNECT " + vsockAddress.getVSOCKPort() + "\n";
      out.write(connectLine.getBytes(StandardCharsets.ISO_8859_1));

      byte[] buf = new byte[16];
      int b;
      int i = 0;
      while ((b = in.read()) != -1 && b != '\n' && i < buf.length) {
        buf[i] = (byte) b;
        if (i < 3) {
          if (OK[i] != b) {
            break;
          }
        }
        i++;
      }
      if (b == '\n' && i > 3) {
        Matcher m = PAT_OK.matcher(new String(buf, 0, i, StandardCharsets.ISO_8859_1));
        if (m.matches()) {
          /* int hostPort = */ Integer.parseInt(m.group(1));
          success = true;
        }
      }
    } finally { // NOPMD.DoNotThrowExceptionInFinally
      if (!success) {
        sock.close();
        throw new SocketException("Unexpected response from proxy socket");
      }
    }

    return sock;
  }
}
