/*	$Id: sockfromto.c,v 1.3 2010/07/09 20:31:48 ryo Exp $	*/

#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>
#include <errno.h>

#include "sockfromto.h"

#undef	PKTINFO_DEBUG
#define RECVFROMTO_WITH_PORT

#if (!defined(IP_PKTINFO) && !defined(IP_SENDSRCADDR)) || !defined(IPV6_PKTINFO)
# warning sockfromto require IP_PKTINFO, IPV6_PKTINFO or IP_SENDSRCADDR
#endif

#ifdef PKTINFO_DEBUG
#define DPRINTF(format, args...)	printf(format, ## args)
#else
#define DPRINTF(args...)
#endif

int
setsockoptfromto(int s)
{
	int rc;
	struct sockaddr_storage ss;
	socklen_t sslen = sizeof(ss);
	int on;

	if ((rc = getsockname(s, (struct sockaddr *)&ss, &sslen)) < 0)
		return rc;

	on = 1;

	switch (((struct sockaddr *)&ss)->sa_family) {
	case AF_INET:
#ifdef IP_RECVDSTADDR
		return setsockopt(s, IPPROTO_IP, IP_RECVDSTADDR, &on, sizeof(on));
#elif defined(IP_RECVPKTINFO)
		return setsockopt(s, IPPROTO_IP, IP_RECVPKTINFO, &on, sizeof(on));
#elif defined(IP_PKTINFO)
		return setsockopt(s, IPPROTO_IP, IP_PKTINFO, &on, sizeof(on));
#else
# warning setsockoptfromto requires IP_RECVDSTADDR, IP_RECVPKTINFO or IP_PKTINFO
#endif

	case AF_INET6:
#ifdef IPV6_RECVPKTINFO
		return setsockopt(s, IPPROTO_IPV6, IPV6_RECVPKTINFO, &on, sizeof(on));
#elif defined(IPV6_PKTINFO)
		return setsockopt(s, IPPROTO_IPV6, IPV6_PKTINFO, &on, sizeof(on));
#else
# warning setsockoptfromto requires IPV6_RECVDSTADDR or IPV6_PKTINFO
#endif

	}

	errno = EAFNOSUPPORT;
	return -1;
}


/*
 *
 *	BSD IPv4
 *			setsockopt(IP_RECVDSTADDR) and
 *			recvmsg with IP_RECVDESTADDR (struct sockaddr_in)
 *
 *	Linux IPv4
 *			setsockopt(IP_PKTINFO) and
 *			recvmsg with IP_PKTINFO (struct ip_pktinfo)
 *
 *	//patched BSD IPv4
 *	//		setsockopt(IP_RECVPKTINFO) and
 *	//		recvmsg with IP_PKTINFO (struct ip_pktinfo)
 *
 *
 *	BSD IPv6 deprecated API (can't use, don't use)
 *			setsockopt(IPV6_RECVDSTADDR) and
 *			recvmsg with IPV6_RECVDESTADDR (struct sockaddr_in6)
 *
 *	BSD IPv6 RFC2292 API (obsolete)
 *	Linux IPv6 RFC2292 API (obsolete)
 *			setsockopt(IPV6_PKTINFO) and
 *			recvmsg with IPV6_PKTINFO (struct in6_pktinfo)
 *
 *	BSD IPv6 RFC2292 compatible API (not recommended)
 *			setsockopt(IPV6_2292PKTINFO) and
 *			recvmsg with IPV6_2292PKTINFO (struct in6_pktinfo)
 *
 *	BSD IPv6 RFC3542 API (recommended)
 *	Linux IPv6 RFC3542 API (recommended)
 *			setsockopt(IPV6_RECVPKTINFO) and
 *			recvmsg with IPV6_PKTINFO (struct in6_pktinfo)
 *
 */
int
recvfromto(int s, void *buf, size_t len, int flags,
		   struct sockaddr *from, socklen_t *fromlen,
		   struct sockaddr *to, socklen_t *tolen)
{
	/* return value */
	int rc;

	socklen_t orig_tolen;
#ifdef RECVFROMTO_WITH_PORT
	/* for getsockname */
	struct sockaddr_storage ss;
	socklen_t sslen = sizeof(ss);
#endif /* RECVFROMTO_WITH_PORT */

	/* for recvmsg */
	struct msghdr msg;
	char cmsgbuf[1024];
	struct iovec vec;

	/* for fetch from cmsghdr */
	struct cmsghdr *cmsg;

	/* fromlen must be a valid pointer */
	if (fromlen == NULL) {
		errno = EFAULT;
		return -1;
	}

	/* setup msghdr */
	vec.iov_base = buf;
	vec.iov_len = len;
	msg.msg_control = cmsgbuf;
	msg.msg_controllen = sizeof(cmsgbuf);
	msg.msg_name = (caddr_t)from;
	msg.msg_namelen = *fromlen;
	msg.msg_iov = &vec;
	msg.msg_iovlen = 1;

	if ((rc = recvmsg(s, &msg, flags)) < 0)
		return rc;

#ifdef RECVFROMTO_WITH_PORT
	if (getsockname(s, (struct sockaddr *)&ss, &sslen) < 0)
		sslen = 0;
#endif

	orig_tolen = *tolen;
	*tolen = 0;
	*fromlen = from->sa_len;

	for (cmsg = (struct cmsghdr *)CMSG_FIRSTHDR(&msg); cmsg;
		cmsg = (struct cmsghdr *)CMSG_NXTHDR(&msg, cmsg)) {

#ifdef IP_RECVDSTADDR
		if ((cmsg->cmsg_level == IPPROTO_IP) && (cmsg->cmsg_type == IP_RECVDSTADDR)) {
			struct in_addr *dstaddr;
			struct sockaddr_in *sin;

			DPRINTF("msg receive: IPPROTO_IP/IP_RECVDESTADDR\n");

			if (orig_tolen < sizeof(struct sockaddr_in)) {
				/* not enough buffer size */
				continue;
			}
			dstaddr = (struct in_addr *)(CMSG_DATA(cmsg));

			*tolen = sizeof(struct sockaddr_in);
			sin = (struct sockaddr_in *)to;
			memset(sin, 0, sizeof(*sin));
			sin->sin_family = AF_INET;
			sin->sin_len = sizeof(*sin);
			sin->sin_addr = *dstaddr;
#ifdef RECVFROMTO_WITH_PORT
			if (sslen)
				sin->sin_port = ((struct sockaddr_in *)&ss)->sin_port;
#endif
			continue;
		}

#elif defined(IP_PKTINFO)
		if ((cmsg->cmsg_level == IPPROTO_IP) && (cmsg->cmsg_type == IP_PKTINFO)) {
			struct in_pktinfo *pktinfo;
			struct sockaddr_in *sin;

			DPRINTF("msg receive: IPPROTO_IP/IP_PKTINFO\n");

			if (orig_tolen < sizeof(struct sockaddr_in)) {
				/* not enough buffer size */
				continue;
			}
			pktinfo = (struct in_pktinfo *)(CMSG_DATA(cmsg));

			*tolen = sizeof(struct sockaddr_in);
			sin = (struct sockaddr_in *)to;
			memset(sin, 0, sizeof(*sin));
			sin->sin_family = AF_INET;
			sin->sin_len = sizeof(*sin);
			sin->sin_addr = pktinfo->ipi_addr;
#ifdef RECVFROMTO_WITH_PORT
			if (sslen)
				sin->sin_port = ((struct sockaddr_in *)&ss)->sin_port;
#endif
			continue;
		}
#endif /* IP_RECVDSTADDR */


#ifdef IPV6_PKTINFO
		if ((cmsg->cmsg_level == IPPROTO_IPV6) && (cmsg->cmsg_type == IPV6_PKTINFO)) {
			struct in6_pktinfo *pktinfo;
			struct sockaddr_in6 *sin6;

			DPRINTF("msg receive: IPPROTO_IPV6/IPV6_PKTINFO\n");

			if (orig_tolen < sizeof(struct sockaddr_in6)) {
				/* not enough buffer size */
				continue;
			}
			pktinfo = (struct in6_pktinfo *)(CMSG_DATA(cmsg));

			*tolen = sizeof(struct sockaddr_in6);
			sin6 = (struct sockaddr_in6 *)to;
			memset(sin6, 0, sizeof(*sin6));
			sin6->sin6_family = AF_INET6;
			sin6->sin6_len = sizeof(*sin6);
			sin6->sin6_addr = pktinfo->ipi6_addr;
			if (IN6_IS_ADDR_LINKLOCAL(&pktinfo->ipi6_addr))
				sin6->sin6_scope_id = pktinfo->ipi6_ifindex;
#ifdef RECVFROMTO_WITH_PORT
			if (sslen)
				sin6->sin6_port = ((struct sockaddr_in6 *)&ss)->sin6_port;
#endif
			continue;
		}
#endif /* IPV6_PKTINFO */

		DPRINTF("%s: unsupported cmsg_level=%d, cmsg_type=%d\n",
			__FUNCTION__,
			cmsg->cmsg_level, cmsg->cmsg_type);
	}

	return rc;
}


/*
 *
 *	*BSD IPv4
 *			socket(), bind() and sendto()
 *
 *	FreeBSD IPv4
 *			sendmsg with IP_SENDSRCADDR (struct in_addr)
 *
 *	//patched BSD IPv4
 *	//		sendmsg with IP_PKTINFO (struct ip_pktinfo)
 *
 *	Linux IPv4
 *			sendmsg with IP_PKTINFO (struct ip_pktinfo)
 *
 *
 *	BSD IPv6 RFC2292/RFC3542 AP (recommended)
 *	Linux IPv6 RFC2292/RFC3542 AP (recommended)
 *			sendmsg with IPV6_PKTINFO (struct in6_pktinfo)
 *
 */
int
sendfromto(int s, const void *buf, size_t len, int flags,
		   const struct sockaddr *from,
		   const struct sockaddr *to)
{
	/* for sendmsg */
	struct msghdr msg;
	char cmsgbuf[1024];
	struct iovec vec;

	/* for store to cmsghdr */
	struct cmsghdr *cmsg;

	/* setup msghdr */
	vec.iov_base = (void *)buf;
	vec.iov_len = len;
	msg.msg_iov = &vec;
	msg.msg_iovlen = 1;
	msg.msg_name = (caddr_t)to;
	msg.msg_namelen = to->sa_len;
	msg.msg_control = cmsgbuf;
	msg.msg_controllen = 0;

	cmsg = (struct cmsghdr *)cmsgbuf;
	switch (to->sa_family) {
	case AF_INET6:
		{
			struct in6_pktinfo *pi;
			int ifindex;

			if (IN6_IS_ADDR_LINKLOCAL(&((const struct sockaddr_in6 *)from)->sin6_addr) ||
				IN6_IS_ADDR_MULTICAST(&((const struct sockaddr_in6 *)from)->sin6_addr))
				ifindex = ((const struct sockaddr_in6 *)from)->sin6_scope_id;
			else
				ifindex = 0;

			msg.msg_controllen += CMSG_SPACE(sizeof(struct in6_pktinfo));
			cmsg->cmsg_len = CMSG_LEN(sizeof(struct in6_pktinfo));
			cmsg->cmsg_level = IPPROTO_IPV6;
			cmsg->cmsg_type = IPV6_PKTINFO;

			pi = (struct in6_pktinfo *)CMSG_DATA(cmsg);
			pi->ipi6_addr = ((const struct sockaddr_in6 *)from)->sin6_addr;
			pi->ipi6_ifindex = ifindex;
		}
		break;
	case AF_INET:
#ifdef IP_PKTINFO
		{
			struct in_pktinfo *pi;

			msg.msg_controllen += CMSG_SPACE(sizeof(struct in_pktinfo));
			cmsg->cmsg_len = CMSG_LEN(sizeof(struct in_pktinfo));
			cmsg->cmsg_level = IPPROTO_IP;
			cmsg->cmsg_type = IP_PKTINFO;

			pi = (struct in_pktinfo *)CMSG_DATA(cmsg);
			pi->ipi_addr = ((const struct sockaddr_in *)from)->sin_addr;
			pi->ipi_ifindex = 0;
		}
#elif defined(IP_SENDSRCADDR)
		{
			struct in_addr *srcadr;

			msg.msg_controllen += CMSG_SPACE(sizeof(struct in_addr));
			cmsg->cmsg_len = CMSG_LEN(sizeof(struct in_addr));
			cmsg->cmsg_level = IPPROTO_IP;
			cmsg->cmsg_type = IP_SENDSRCADDR;

			srcadr = (struct in_addr*)CMSG_DATA(cmsg);
			*srcadr = ((const struct sockaddr_in *)from)->sin_addr;
		}
#else
# warning no IP_PKTINFO nor IP_SENDSRCADDR
#endif
		break;
	}

	if (msg.msg_controllen == 0)
		msg.msg_control = NULL;

	return sendmsg(s, &msg, flags);
}


/*
 * omake
 */
char *
strsockaddr(struct sockaddr *sa)
{
#define NBUF	8
#define BUFSIZE 256
	static int n = 0;
	static char buf[NBUF][BUFSIZE];
	char *p0, *p;

	p0 = p = buf[n++ % NBUF];
	n %= NBUF;

	switch (sa->sa_family) {
	case AF_INET:
		inet_ntop(AF_INET, &((struct sockaddr_in *)sa)->sin_addr, p, BUFSIZE);
		p = p + strlen(p);
		sprintf(p, ":%d", ntohs(((struct sockaddr_in *)sa)->sin_port));
		break;
	case AF_INET6:
		*p++ = '[';
		inet_ntop(AF_INET6, &((struct sockaddr_in6 *)sa)->sin6_addr, p, BUFSIZE);
		p = p + strlen(p);
		sprintf(p, "]:%d", ntohs(((struct sockaddr_in6 *)sa)->sin6_port));
		break;
	default:
		{
			unsigned char *s = (unsigned char *)sa->sa_data;
			unsigned char *end = (unsigned char *)sa + sa->sa_len;
			int left;
			p += sprintf(p, "(%d)", sa->sa_family);

			for (left = BUFSIZE; s < end && left > 2; left -= 2)
				p += snprintf(p, left, "%02x", *s++);
		}
		break;
	}

	return p0;
}