// SPDX-License-Identifier: Apache-2.0
//
// AF_INET sockets, on top of lwIP's raw API.
//
// The raw API rather than lwIP's own sockets layer, because lwIP's needs
// NO_SYS=0 -- threads, mailboxes and semaphores of its own -- and this kernel
// already has a socket layer with a shape of its own that AF_INET has to fit
// into rather than replace. See kernel/net/lwip/lwipopts.h.
//
// ---------------------------------------------------------------------------
// Locking, and why there is only one lock here
//
// Everything in an inet socket is protected by the lwIP lock and nothing else.
// That is not a shortcut: lwIP's callbacks (data arriving, a connection
// completing, a peer closing) run in the network thread with that lock held,
// and they are the only things that write to a socket's receive buffer. A
// syscall reading that buffer takes the same lock. There is no second lock
// because there is no second writer.
//
// ---------------------------------------------------------------------------
// Blocking, and what it costs today
//
// A blocking read waits by polling on the timer tick rather than by sleeping on
// a channel a callback wakes. The reason is a lock mismatch: sleep() wants a
// spinlock, the lwIP lock is a sleeplock (it must be, because a call into lwIP
// can reach the driver and the driver can sleep), and interleaving the two
// correctly is a genuinely error-prone piece of code.
//
// The cost is up to 10 ms of latency on an operation that has to wait -- and
// none at all on one that does not, which is the case that matters for
// throughput: a bulk transfer only blocks when the buffer is genuinely empty.
// Replacing this with a proper sleep/wakeup is worth doing and is deliberately
// not the first version.

#include "types.h"
#include "defs.h"
#include "param.h"
#include "spinlock.h"
#include "sleeplock.h"
#include "proc.h"
#include "fs.h"      /* NDIRECT, which file.h's struct inode needs */
#include "file.h"
#define _KERNEL_SOCKET_INTERNALS
#include "socket.h"
#include "errno.h"

#include "lwip/opt.h"
#include "lwip/tcp.h"
#include "lwip/udp.h"
#include "lwip/ip_addr.h"
#include "lwip/pbuf.h"

#define NINET      16        // AF_INET sockets in the system
#define INET_RXBUF 8192      // per socket; one lwIP window plus slack

struct inet_sock {
  int    inuse;
  int    type;               // SOCK_STREAM or SOCK_DGRAM

  struct tcp_pcb *tpcb;
  struct udp_pcb *upcb;

  // Receive ring. lwIP hands over pbufs from the network thread; a reader
  // drains bytes. Full means the sender is told to back off, which for TCP is
  // the window closing rather than data being dropped.
  char   rx[INET_RXBUF];
  uint   rhead, rtail;       // rhead == rtail is empty

  int    connecting;         // connect() issued, no answer yet
  int    connected;
  int    conn_failed;
  int    peer_closed;        // FIN seen, or the connection died
  int    listening;

  // Connections lwIP has completed but nobody has accept()ed yet.
  struct inet_sock *acceptq[NBACKLOG];
  int    naccept;
};

static struct inet_sock itab[NINET];

/* ------------------------------------------------------------------ */
/* The receive ring. Callers hold the lwIP lock.                       */
/* ------------------------------------------------------------------ */

static uint
rx_used(struct inet_sock *is)
{
  return is->rhead - is->rtail;
}

static uint
rx_space(struct inet_sock *is)
{
  return INET_RXBUF - rx_used(is);
}

static void
rx_put(struct inet_sock *is, const char *p, uint n)
{
  for (uint i = 0; i < n; i++)
    is->rx[(is->rhead + i) % INET_RXBUF] = p[i];
  is->rhead += n;
}

static uint
rx_take(struct inet_sock *is, char *p, uint n)
{
  uint have = rx_used(is);
  if (n > have)
    n = have;
  for (uint i = 0; i < n; i++)
    p[i] = is->rx[(is->rtail + i) % INET_RXBUF];
  is->rtail += n;
  return n;
}

/* ------------------------------------------------------------------ */
/* lwIP callbacks. All of these run in the network thread, holding the */
/* lwIP lock.                                                          */
/* ------------------------------------------------------------------ */

static void
tcp_fail(struct inet_sock *is)
{
  is->peer_closed = 1;
  is->conn_failed = 1;
  is->connecting  = 0;
  // The pcb is already freed by lwIP when err_cb fires; keeping the pointer
  // would leave a dangling handle that close() would then free a second time.
  is->tpcb = 0;
}

static void
inet_err_cb(void *arg, err_t err)
{
  (void)err;
  if (arg)
    tcp_fail((struct inet_sock *)arg);
}

static err_t
inet_recv_cb(void *arg, struct tcp_pcb *pcb, struct pbuf *p, err_t err)
{
  struct inet_sock *is = arg;

  if (is == 0) {
    if (p) pbuf_free(p);
    return ERR_OK;
  }
  if (err != ERR_OK) {
    if (p) pbuf_free(p);
    return err;
  }
  // A null pbuf is the peer's FIN, not an error: the connection stays readable
  // until the buffer drains, and only then does read() report end of file.
  if (p == NULL) {
    is->peer_closed = 1;
    return ERR_OK;
  }

  // Refusing what will not fit leaves it queued in lwIP, which closes the
  // window and makes the sender wait -- which is the whole point of a window.
  // Accepting and dropping would look like a working receive that loses data.
  if (p->tot_len > rx_space(is))
    return ERR_MEM;

  for (struct pbuf *q = p; q != NULL; q = q->next)
    rx_put(is, (const char *)q->payload, q->len);

  // Only now is the data ours, so only now may the window reopen.
  tcp_recved(pcb, p->tot_len);
  pbuf_free(p);
  return ERR_OK;
}

static err_t
inet_connected_cb(void *arg, struct tcp_pcb *pcb, err_t err)
{
  struct inet_sock *is = arg;
  (void)pcb;

  if (is == 0)
    return ERR_OK;
  if (err != ERR_OK) {
    tcp_fail(is);
    return ERR_OK;
  }
  is->connecting = 0;
  is->connected  = 1;
  return ERR_OK;
}

static struct inet_sock *inet_alloc(int type);

static err_t
inet_accept_cb(void *arg, struct tcp_pcb *newpcb, err_t err)
{
  struct inet_sock *ls = arg;

  if (ls == 0 || err != ERR_OK)
    return ERR_VAL;
  if (ls->naccept >= NBACKLOG)
    return ERR_MEM;            // backlog full: lwIP refuses the connection

  struct inet_sock *is = inet_alloc(SOCK_STREAM);
  if (is == 0)
    return ERR_MEM;

  is->tpcb      = newpcb;
  is->connected = 1;
  tcp_arg(newpcb, is);
  tcp_recv(newpcb, inet_recv_cb);
  tcp_err(newpcb, inet_err_cb);

  ls->acceptq[ls->naccept++] = is;
  return ERR_OK;
}

// UDP delivers datagrams with no connection. The source is dropped on the
// floor for now: recvfrom is not wired, and inventing a per-datagram source
// queue that nothing reads would be storage with no reader.
static void
inet_udp_recv_cb(void *arg, struct udp_pcb *pcb, struct pbuf *p,
                 const ip_addr_t *addr, u16_t port)
{
  struct inet_sock *is = arg;
  (void)pcb; (void)addr; (void)port;

  if (is == 0 || p == NULL) {
    if (p) pbuf_free(p);
    return;
  }
  if (p->tot_len <= rx_space(is))
    for (struct pbuf *q = p; q != NULL; q = q->next)
      rx_put(is, (const char *)q->payload, q->len);
  pbuf_free(p);
}

/* ------------------------------------------------------------------ */
/* Allocation                                                          */
/* ------------------------------------------------------------------ */

static struct inet_sock *
inet_alloc(int type)
{
  for (int i = 0; i < NINET; i++)
    if (!itab[i].inuse) {
      struct inet_sock *is = &itab[i];
      memset(is, 0, sizeof(*is));
      is->inuse = 1;
      is->type  = type;
      return is;
    }
  return 0;
}

// Wait one timer tick. See the note at the top on why waiting is a poll.
static void
inet_wait(void)
{
  acquire(&tickslock);
  uint start = ticks;
  while (ticks == start)
    sleep(&ticks, &tickslock);
  release(&tickslock);
}

/* ------------------------------------------------------------------ */
/* The interface kernel/net/unix.c dispatches to                       */
/* ------------------------------------------------------------------ */

void *
inet_new(int type)
{
  if (type != SOCK_STREAM && type != SOCK_DGRAM)
    return 0;
  if (!net_is_up())
    return 0;

  lwip_sys_lock();
  struct inet_sock *is = inet_alloc(type);
  if (is == 0) {
    lwip_sys_unlock();
    return 0;
  }

  if (type == SOCK_STREAM) {
    is->tpcb = tcp_new();
    if (is->tpcb == 0) {
      is->inuse = 0;
      lwip_sys_unlock();
      return 0;
    }
    tcp_arg(is->tpcb, is);
    tcp_err(is->tpcb, inet_err_cb);
    tcp_recv(is->tpcb, inet_recv_cb);
  } else {
    is->upcb = udp_new();
    if (is->upcb == 0) {
      is->inuse = 0;
      lwip_sys_unlock();
      return 0;
    }
    udp_recv(is->upcb, inet_udp_recv_cb, is);
  }

  lwip_sys_unlock();
  return is;
}

int
inet_bind(void *h, uint32 addr, uint16 port)
{
  struct inet_sock *is = h;
  ip_addr_t ip;
  err_t e;

  ip_addr_set_ip4_u32(&ip, addr);      // already network byte order

  lwip_sys_lock();
  e = (is->type == SOCK_STREAM) ? tcp_bind(is->tpcb, &ip, port)
                                : udp_bind(is->upcb, &ip, port);
  lwip_sys_unlock();

  if (e == ERR_USE)
    return -EADDRINUSE;
  return e == ERR_OK ? 0 : -EINVAL;
}

int
inet_listen(void *h, int backlog)
{
  struct inet_sock *is = h;

  if (is->type != SOCK_STREAM)
    return -EOPNOTSUPP;

  lwip_sys_lock();
  // tcp_listen frees the old pcb and returns a smaller listening one, so the
  // handle has to be replaced rather than kept. Losing this is a use-after-free
  // that only shows up on the second connection.
  struct tcp_pcb *lp = tcp_listen_with_backlog(is->tpcb,
                          backlog > 0 && backlog < NBACKLOG ? (u8_t)backlog
                                                            : NBACKLOG);
  if (lp == 0) {
    lwip_sys_unlock();
    return -ENOMEM;
  }
  is->tpcb      = lp;
  is->listening = 1;
  tcp_arg(lp, is);
  tcp_accept(lp, inet_accept_cb);
  lwip_sys_unlock();
  return 0;
}

void *
inet_accept(void *h)
{
  struct inet_sock *ls = h;

  if (!ls->listening)
    return 0;

  for (;;) {
    lwip_sys_lock();
    if (ls->naccept > 0) {
      struct inet_sock *is = ls->acceptq[0];
      for (int i = 1; i < ls->naccept; i++)
        ls->acceptq[i - 1] = ls->acceptq[i];
      ls->naccept--;
      lwip_sys_unlock();
      return is;
    }
    lwip_sys_unlock();

    if (myproc()->killed)
      return 0;
    inet_wait();
  }
}

int
inet_connect(void *h, uint32 addr, uint16 port)
{
  struct inet_sock *is = h;
  ip_addr_t ip;

  ip_addr_set_ip4_u32(&ip, addr);

  if (is->type == SOCK_DGRAM) {
    // A connected UDP socket just remembers where to send; there is no
    // handshake to wait for.
    lwip_sys_lock();
    err_t e = udp_connect(is->upcb, &ip, port);
    lwip_sys_unlock();
    return e == ERR_OK ? 0 : -EINVAL;
  }

  lwip_sys_lock();
  is->connecting = 1;
  err_t e = tcp_connect(is->tpcb, &ip, port, inet_connected_cb);
  lwip_sys_unlock();
  if (e != ERR_OK) {
    is->connecting = 0;
    return -ECONNREFUSED;
  }

  for (;;) {
    lwip_sys_lock();
    int done = is->connected, failed = is->conn_failed;
    lwip_sys_unlock();
    if (done)
      return 0;
    if (failed)
      return -ECONNREFUSED;
    if (myproc()->killed)
      return -EINTR;
    inet_wait();
  }
}

int
inet_read(void *h, char *buf, int n)
{
  struct inet_sock *is = h;

  if (n <= 0)
    return 0;

  for (;;) {
    lwip_sys_lock();
    uint got = rx_take(is, buf, (uint)n);
    int  eof = is->peer_closed;
    if (got > 0 && is->type == SOCK_STREAM && is->tpcb)
      ;  // window already reopened in the receive callback
    lwip_sys_unlock();

    if (got > 0)
      return (int)got;
    if (eof)
      return 0;                       // end of file, not an error
    if (myproc()->killed)
      return -EINTR;
    inet_wait();
  }
}

int
inet_write(void *h, char *buf, int n)
{
  struct inet_sock *is = h;
  int done = 0;

  if (n <= 0)
    return 0;

  if (is->type == SOCK_DGRAM) {
    lwip_sys_lock();
    struct pbuf *p = pbuf_alloc(PBUF_TRANSPORT, (u16_t)n, PBUF_RAM);
    if (p == 0) {
      lwip_sys_unlock();
      return -ENOBUFS;
    }
    pbuf_take(p, buf, (u16_t)n);
    err_t e = udp_send(is->upcb, p);
    pbuf_free(p);
    lwip_sys_unlock();
    return e == ERR_OK ? n : -EIO;
  }

  // TCP: the send buffer is finite, so a large write becomes several passes.
  while (done < n) {
    lwip_sys_lock();
    if (is->peer_closed || is->tpcb == 0) {
      lwip_sys_unlock();
      // Same contract as a pipe with no reader, which is what the AF_UNIX side
      // does: the signal first, then the error.
      return done > 0 ? done : -EPIPE;
    }

    int room = tcp_sndbuf(is->tpcb);
    int want = n - done;
    if (want > room)
      want = room;

    if (want > 0) {
      // TCP_WRITE_FLAG_COPY because `buf` is a kernel bounce buffer the caller
      // reuses the moment this returns; lwIP would otherwise hold a pointer
      // into memory that has changed by the time the segment goes out.
      err_t e = tcp_write(is->tpcb, buf + done, (u16_t)want,
                          TCP_WRITE_FLAG_COPY);
      if (e == ERR_OK) {
        tcp_output(is->tpcb);
        done += want;
      } else if (e != ERR_MEM) {
        lwip_sys_unlock();
        return done > 0 ? done : -EIO;
      }
    }
    lwip_sys_unlock();

    if (done < n) {
      if (myproc()->killed)
        return done > 0 ? done : -EINTR;
      inet_wait();                    // wait for the window or the buffer
    }
  }
  return done;
}

int
inet_readable(void *h)
{
  struct inet_sock *is = h;
  int r;

  lwip_sys_lock();
  r = rx_used(is) > 0 || is->peer_closed || is->naccept > 0;
  lwip_sys_unlock();
  return r;
}

int
inet_hup(void *h)
{
  struct inet_sock *is = h;
  int r;

  lwip_sys_lock();
  r = is->peer_closed && rx_used(is) == 0;
  lwip_sys_unlock();
  return r;
}

void
inet_close(void *h)
{
  struct inet_sock *is = h;

  if (is == 0)
    return;

  lwip_sys_lock();
  // Anything queued but never accepted is closed too, or those connections
  // would sit half-open until the peer gave up on them.
  for (int i = 0; i < is->naccept; i++) {
    struct inet_sock *q = is->acceptq[i];
    if (q->tpcb) {
      tcp_arg(q->tpcb, 0);
      tcp_close(q->tpcb);
    }
    q->inuse = 0;
  }
  is->naccept = 0;

  if (is->tpcb) {
    // Clearing the argument first: lwIP may still deliver a callback for this
    // pcb after tcp_close, and it would otherwise arrive with a pointer to a
    // socket that has just been freed.
    tcp_arg(is->tpcb, 0);
    tcp_recv(is->tpcb, 0);
    tcp_err(is->tpcb, 0);
    if (tcp_close(is->tpcb) != ERR_OK)
      tcp_abort(is->tpcb);
    is->tpcb = 0;
  }
  if (is->upcb) {
    udp_recv(is->upcb, 0, 0);
    udp_remove(is->upcb);
    is->upcb = 0;
  }
  is->inuse = 0;
  lwip_sys_unlock();
}