/* -*- mode: c; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 * Copyright 2016 by the Massachusetts Institute of Technology.
 * All Rights Reserved.
 *
 * Export of this software from the United States of America may
 * require a specific license from the United States Government.
 * It is the responsibility of any person or organization contemplating
 * export to obtain such a license before exporting.
 *
 * WITHIN THAT CONSTRAINT, permission to use, copy, modify, and
 * distribute this software and its documentation for any purpose and
 * without fee is hereby granted, provided that the above copyright
 * notice appear in all copies and that both that copyright notice and
 * this permission notice appear in supporting documentation, and that
 * the name of M.I.T. not be used in advertising or publicity pertaining
 * to distribution of the software without specific, written prior
 * permission.  Furthermore if you modify this software you must label
 * your software as modified software and not distribute it in such a
 * fashion that it might be confused with the original M.I.T. software.
 * M.I.T. makes no representations about the suitability of
 * this software for any purpose.  It is provided "as is" without express
 * or implied warranty.
 */

/* macOS requires this define for IPV6_PKTINFO. */
#define __APPLE_USE_RFC_3542

#include "udppktinfo.h"

#include <netinet/in.h>
#include <sys/socket.h>

#if defined(IP_PKTINFO) && defined(HAVE_STRUCT_IN_PKTINFO)
#define HAVE_IP_PKTINFO
#endif

#if defined(IPV6_PKTINFO) && defined(HAVE_STRUCT_IN6_PKTINFO)
#define HAVE_IPV6_PKTINFO
#endif

#if defined(HAVE_IP_PKTINFO) || defined(IP_SENDSRCADDR) ||      \
    defined(HAVE_IPV6_PKTINFO)
#define HAVE_PKTINFO_SUPPORT
#endif

/* Use RFC 3542 API below, but fall back from IPV6_RECVPKTINFO to IPV6_PKTINFO
 * for RFC 2292 implementations. */
#if !defined(IPV6_RECVPKTINFO) && defined(IPV6_PKTINFO)
#define IPV6_RECVPKTINFO IPV6_PKTINFO
#endif

/* Parallel, though not standardized. */
#if !defined(IP_RECVPKTINFO) && defined(IP_PKTINFO)
#define IP_RECVPKTINFO IP_PKTINFO
#endif /* IP_RECVPKTINFO */

#if defined(CMSG_SPACE) && defined(HAVE_STRUCT_CMSGHDR) &&      \
    defined(HAVE_PKTINFO_SUPPORT)
union pktinfo {
#ifdef HAVE_STRUCT_IN6_PKTINFO
    struct in6_pktinfo pi6;
#endif
#ifdef HAVE_STRUCT_IN_PKTINFO
    struct in_pktinfo pi4;
#endif
#ifdef IP_RECVDSTADDR
    struct in_addr iaddr;
#endif
    char c;
};
#endif /* HAVE_IPV6_PKTINFO && HAVE_STRUCT_CMSGHDR && HAVE_PKTINFO_SUPPORT */

#ifdef HAVE_IP_PKTINFO

#define set_ipv4_pktinfo set_ipv4_recvpktinfo
static inline krb5_error_code
set_ipv4_recvpktinfo(int sock)
{
    int sockopt = 1;
    return setsockopt(sock, IPPROTO_IP, IP_RECVPKTINFO, &sockopt,
                      sizeof(sockopt));
}

#elif defined(IP_RECVDSTADDR) /* HAVE_IP_PKTINFO */

#define set_ipv4_pktinfo set_ipv4_recvdstaddr
static inline krb5_error_code
set_ipv4_recvdstaddr(int sock)
{
    int sockopt = 1;
    return setsockopt(sock, IPPROTO_IP, IP_RECVDSTADDR, &sockopt,
                      sizeof(sockopt));
}

#else /* HAVE_IP_PKTINFO || IP_RECVDSTADDR */
#define set_ipv4_pktinfo(s) EINVAL
#endif /* HAVE_IP_PKTINFO || IP_RECVDSTADDR */

#ifdef HAVE_IPV6_PKTINFO

#define set_ipv6_pktinfo set_ipv6_recvpktinfo
static inline krb5_error_code
set_ipv6_recvpktinfo(int sock)
{
    int sockopt = 1;
    return setsockopt(sock, IPPROTO_IPV6, IPV6_RECVPKTINFO, &sockopt,
                      sizeof(sockopt));
}

#else /* HAVE_IPV6_PKTINFO */
#define set_ipv6_pktinfo(s) EINVAL
#endif /* HAVE_IPV6_PKTINFO */

/*
 * Set pktinfo option on a socket. Takes a socket and the socket address family
 * as arguments.
 *
 * Returns 0 on success, EINVAL if pktinfo is not supported for the address
 * family.
 */
krb5_error_code
set_pktinfo(int sock, int family)
{
    switch (family) {
    case AF_INET:
        return set_ipv4_pktinfo(sock);
    case AF_INET6:
        return set_ipv6_pktinfo(sock);
    default:
        return EINVAL;
    }
}

#if defined(HAVE_PKTINFO_SUPPORT) && defined(CMSG_SPACE)

/*
 * Check if a socket is bound to a wildcard address.
 * Returns 1 if it is, 0 if it's bound to a specific address, or -1 on error
 * with errno set to the error.
 */
static int
is_socket_bound_to_wildcard(int sock)
{
    struct sockaddr_storage bound_addr;
    socklen_t bound_addr_len = sizeof(bound_addr);
    struct sockaddr *sa = ss2sa(&bound_addr);

    if (getsockname(sock, sa, &bound_addr_len) < 0)
        return -1;

    if (!sa_is_inet(sa)) {
        errno = EINVAL;
        return -1;
    }

    return sa_is_wildcard(sa);
}

#ifdef HAVE_IP_PKTINFO

static inline struct in_pktinfo *
cmsg2pktinfo(struct cmsghdr *cmsgptr)
{
    return (struct in_pktinfo *)(void *)CMSG_DATA(cmsgptr);
}

#define check_cmsg_v4_pktinfo check_cmsg_ip_pktinfo
static int
check_cmsg_ip_pktinfo(struct cmsghdr *cmsgptr, struct sockaddr *to,
                      socklen_t *tolen, aux_addressing_info *auxaddr)
{
    struct in_pktinfo *pktinfo;

    if (cmsgptr->cmsg_level == IPPROTO_IP &&
        cmsgptr->cmsg_type == IP_PKTINFO &&
        *tolen >= sizeof(struct sockaddr_in)) {

        memset(to, 0, sizeof(struct sockaddr_in));
        pktinfo = cmsg2pktinfo(cmsgptr);
        sa2sin(to)->sin_addr = pktinfo->ipi_addr;
        sa2sin(to)->sin_family = AF_INET;
        *tolen = sizeof(struct sockaddr_in);
        return 1;
    }
    return 0;
}

#elif defined(IP_RECVDSTADDR) /* HAVE_IP_PKTINFO */

static inline struct in_addr *
cmsg2sin(struct cmsghdr *cmsgptr)
{
    return (struct in_addr *)(void *)CMSG_DATA(cmsgptr);
}

#define check_cmsg_v4_pktinfo check_cmsg_ip_recvdstaddr
static int
check_cmsg_ip_recvdstaddr(struct cmsghdr *cmsgptr, struct sockaddr *to,
                          socklen_t *tolen, aux_addressing_info * auxaddr)
{
    if (cmsgptr->cmsg_level == IPPROTO_IP &&
        cmsgptr->cmsg_type == IP_RECVDSTADDR &&
        *tolen >= sizeof(struct sockaddr_in)) {
        struct in_addr *sin_addr;

        memset(to, 0, sizeof(struct sockaddr_in));
        sin_addr = cmsg2sin(cmsgptr);
        sa2sin(to)->sin_addr = *sin_addr;
        sa2sin(to)->sin_family = AF_INET;
        *tolen = sizeof(struct sockaddr_in);
        return 1;
    }
    return 0;
}

#else /* HAVE_IP_PKTINFO || IP_RECVDSTADDR */
#define check_cmsg_v4_pktinfo(c, t, l, a) 0
#endif /* HAVE_IP_PKTINFO || IP_RECVDSTADDR */

#ifdef HAVE_IPV6_PKTINFO

static inline struct in6_pktinfo *
cmsg2pktinfo6(struct cmsghdr *cmsgptr)
{
    return (struct in6_pktinfo *)(void *)CMSG_DATA(cmsgptr);
}

#define check_cmsg_v6_pktinfo check_cmsg_ipv6_pktinfo
static int
check_cmsg_ipv6_pktinfo(struct cmsghdr *cmsgptr, struct sockaddr *to,
                        socklen_t *tolen, aux_addressing_info *auxaddr)
{
    struct in6_pktinfo *pktinfo;

    if (cmsgptr->cmsg_level == IPPROTO_IPV6 &&
        cmsgptr->cmsg_type == IPV6_PKTINFO &&
        *tolen >= sizeof(struct sockaddr_in6)) {

        memset(to, 0, sizeof(struct sockaddr_in6));
        pktinfo = cmsg2pktinfo6(cmsgptr);
        sa2sin6(to)->sin6_addr = pktinfo->ipi6_addr;
        sa2sin6(to)->sin6_family = AF_INET6;
        *tolen = sizeof(struct sockaddr_in6);
        auxaddr->ipv6_ifindex = pktinfo->ipi6_ifindex;
        return 1;
    }
    return 0;
}
#else /* HAVE_IPV6_PKTINFO */
#define check_cmsg_v6_pktinfo(c, t, l, a) 0
#endif /* HAVE_IPV6_PKTINFO */

static int
check_cmsg_pktinfo(struct cmsghdr *cmsgptr, struct sockaddr *to,
                   socklen_t *tolen, aux_addressing_info *auxaddr)
{
    return check_cmsg_v4_pktinfo(cmsgptr, to, tolen, auxaddr) ||
           check_cmsg_v6_pktinfo(cmsgptr, to, tolen, auxaddr);
}

/*
 * Receive a message from a socket.
 *
 * Arguments:
 *  sock
 *  buf     - The buffer to store the message in.
 *  len     - buf length
 *  flags
 *  from    - Set to the address that sent the message
 *  fromlen
 *  to      - Set to the address that the message was sent to if possible.
 *            May not be set in certain cases such as if pktinfo support is
 *            missing. May be NULL.
 *  tolen
 *  auxaddr - Miscellaneous address information.
 *
 * Returns 0 on success, otherwise an error code.
 */
krb5_error_code
recv_from_to(int sock, void *buf, size_t len, int flags,
             struct sockaddr *from, socklen_t * fromlen,
             struct sockaddr *to, socklen_t * tolen,
             aux_addressing_info *auxaddr)

{
    int r;
    struct iovec iov;
    char cmsg[CMSG_SPACE(sizeof(union pktinfo))];
    struct cmsghdr *cmsgptr;
    struct msghdr msg;

    /* Don't use pktinfo if the socket isn't bound to a wildcard address. */
    r = is_socket_bound_to_wildcard(sock);
    if (r < 0)
        return errno;

    if (!to || !tolen || !r)
        return recvfrom(sock, buf, len, flags, from, fromlen);

    /* Clobber with something recognizeable in case we can't extract the
     * address but try to use it anyways. */
    memset(to, 0x40, *tolen);

    iov.iov_base = buf;
    iov.iov_len = len;
    memset(&msg, 0, sizeof(msg));
    msg.msg_name = from;
    msg.msg_namelen = *fromlen;
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = cmsg;
    msg.msg_controllen = sizeof(cmsg);

    r = recvmsg(sock, &msg, flags);
    if (r < 0)
        return r;
    *fromlen = msg.msg_namelen;

    /*
     * On Darwin (and presumably all *BSD with KAME stacks), CMSG_FIRSTHDR
     * doesn't check for a non-zero controllen.  RFC 3542 recommends making
     * this check, even though the (new) spec for CMSG_FIRSTHDR says it's
     * supposed to do the check.
     */
    if (msg.msg_controllen) {
        cmsgptr = CMSG_FIRSTHDR(&msg);
        while (cmsgptr) {
            if (check_cmsg_pktinfo(cmsgptr, to, tolen, auxaddr))
                return r;
            cmsgptr = CMSG_NXTHDR(&msg, cmsgptr);
        }
    }
    /* No info about destination addr was available.  */
    *tolen = 0;
    return r;
}

#ifdef HAVE_IP_PKTINFO

#define set_msg_from_ipv4 set_msg_from_ip_pktinfo
static krb5_error_code
set_msg_from_ip_pktinfo(struct msghdr *msg, struct cmsghdr *cmsgptr,
                        struct sockaddr *from, socklen_t fromlen,
                        aux_addressing_info *auxaddr)
{
    struct in_pktinfo *p = cmsg2pktinfo(cmsgptr);
    const struct sockaddr_in *from4 = sa2sin(from);

    if (fromlen != sizeof(struct sockaddr_in))
        return EINVAL;
    cmsgptr->cmsg_level = IPPROTO_IP;
    cmsgptr->cmsg_type = IP_PKTINFO;
    cmsgptr->cmsg_len = CMSG_LEN(sizeof(struct in_pktinfo));
    p->ipi_spec_dst = from4->sin_addr;

    msg->msg_controllen = CMSG_SPACE(sizeof(struct in_pktinfo));
    return 0;
}

#elif defined(IP_SENDSRCADDR) /* HAVE_IP_PKTINFO */

#define set_msg_from_ipv4 set_msg_from_ip_sendsrcaddr
static krb5_error_code
set_msg_from_ip_sendsrcaddr(struct msghdr *msg, struct cmsghdr *cmsgptr,
                            struct sockaddr *from, socklen_t fromlen,
                            aux_addressing_info *auxaddr)
{
    struct in_addr *sin_addr = cmsg2sin(cmsgptr);
    const struct sockaddr_in *from4 = sa2sin(from);
    if (fromlen != sizeof(struct sockaddr_in))
        return EINVAL;
    cmsgptr->cmsg_level = IPPROTO_IP;
    cmsgptr->cmsg_type = IP_SENDSRCADDR;
    cmsgptr->cmsg_len = CMSG_LEN(sizeof(struct in_addr));
    msg->msg_controllen = CMSG_SPACE(sizeof(struct in_addr));
    *sin_addr = from4->sin_addr;
    return 0;
}

#else /* HAVE_IP_PKTINFO || IP_SENDSRCADDR */
#define set_msg_from_ipv4(m, c, f, l, a) EINVAL
#endif /* HAVE_IP_PKTINFO || IP_SENDSRCADDR */

#ifdef HAVE_IPV6_PKTINFO

#define set_msg_from_ipv6 set_msg_from_ipv6_pktinfo
static krb5_error_code
set_msg_from_ipv6_pktinfo(struct msghdr *msg, struct cmsghdr *cmsgptr,
                          struct sockaddr *from, socklen_t fromlen,
                          aux_addressing_info *auxaddr)
{
    struct in6_pktinfo *p = cmsg2pktinfo6(cmsgptr);
    const struct sockaddr_in6 *from6 = sa2sin6(from);

    if (fromlen != sizeof(struct sockaddr_in6))
        return EINVAL;
    cmsgptr->cmsg_level = IPPROTO_IPV6;
    cmsgptr->cmsg_type = IPV6_PKTINFO;
    cmsgptr->cmsg_len = CMSG_LEN(sizeof(struct in6_pktinfo));

    p->ipi6_addr = from6->sin6_addr;
    /*
     * Because of the possibility of asymmetric routing, we
     * normally don't want to specify an interface.  However,
     * macOS doesn't like sending from a link-local address
     * (which can come up in testing at least, if you wind up
     * with a "foo.local" name) unless we do specify the
     * interface.
     */
    if (IN6_IS_ADDR_LINKLOCAL(&from6->sin6_addr))
        p->ipi6_ifindex = auxaddr->ipv6_ifindex;
    /* otherwise, already zero */

    msg->msg_controllen = CMSG_SPACE(sizeof(struct in6_pktinfo));
    return 0;
}

#else /* HAVE_IPV6_PKTINFO */
#define set_msg_from_ipv6(m, c, f, l, a) EINVAL
#endif /* HAVE_IPV6_PKTINFO */

static krb5_error_code
set_msg_from(int family, struct msghdr *msg, struct cmsghdr *cmsgptr,
             struct sockaddr *from, socklen_t fromlen,
             aux_addressing_info *auxaddr)
{
    switch (family) {
    case AF_INET:
        return set_msg_from_ipv4(msg, cmsgptr, from, fromlen, auxaddr);
    case AF_INET6:
        return set_msg_from_ipv6(msg, cmsgptr, from, fromlen, auxaddr);
    }

    return EINVAL;
}

/*
 * Send a message to an address.
 *
 * Arguments:
 *  sock
 *  buf     - The message to send.
 *  len     - buf length
 *  flags
 *  to      - The address to send the message to.
 *  tolen
 *  from    - The address to attempt to send the message from. May be NULL.
 *  fromlen
 *  auxaddr - Miscellaneous address information.
 *
 * Returns 0 on success, otherwise an error code.
 */
krb5_error_code
send_to_from(int sock, void *buf, size_t len, int flags,
             const struct sockaddr *to, socklen_t tolen, struct sockaddr *from,
             socklen_t fromlen, aux_addressing_info *auxaddr)
{
    int r;
    struct iovec iov;
    struct msghdr msg;
    struct cmsghdr *cmsgptr;
    char cbuf[CMSG_SPACE(sizeof(union pktinfo))];

    /* Don't use pktinfo if the socket isn't bound to a wildcard address. */
    r = is_socket_bound_to_wildcard(sock);
    if (r < 0)
        return errno;

    if (from == NULL || fromlen == 0 || from->sa_family != to->sa_family || !r)
        goto use_sendto;

    iov.iov_base = buf;
    iov.iov_len = len;
    /* Truncation?  */
    if (iov.iov_len != len)
        return EINVAL;
    memset(cbuf, 0, sizeof(cbuf));
    memset(&msg, 0, sizeof(msg));
    msg.msg_name = (void *)to;
    msg.msg_namelen = tolen;
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = cbuf;
    /* CMSG_FIRSTHDR needs a non-zero controllen, or it'll return NULL on
     * Linux. */
    msg.msg_controllen = sizeof(cbuf);
    cmsgptr = CMSG_FIRSTHDR(&msg);
    msg.msg_controllen = 0;

    if (set_msg_from(from->sa_family, &msg, cmsgptr, from, fromlen, auxaddr))
        goto use_sendto;
    return sendmsg(sock, &msg, flags);

use_sendto:
    return sendto(sock, buf, len, flags, to, tolen);
}

#else /* HAVE_PKTINFO_SUPPORT && CMSG_SPACE */

krb5_error_code
recv_from_to(int sock, void *buf, size_t len, int flags,
             struct sockaddr *from, socklen_t *fromlen,
             struct sockaddr *to, socklen_t *tolen,
             aux_addressing_info *auxaddr)
{
    if (to && tolen) {
        /* Clobber with something recognizeable in case we try to use the
         * address. */
        memset(to, 0x40, *tolen);
        *tolen = 0;
    }

    return recvfrom(sock, buf, len, flags, from, fromlen);
}

krb5_error_code
send_to_from(int sock, void *buf, size_t len, int flags,
             const struct sockaddr *to, socklen_t tolen,
             struct sockaddr *from, socklen_t fromlen,
             aux_addressing_info *auxaddr)
{
    return sendto(sock, buf, len, flags, to, tolen);
}

#endif /* HAVE_PKTINFO_SUPPORT && CMSG_SPACE */