/*
  gcc -fPIC -shared sendmsg.c -o sendmsg.so -ldl
  LD_PRELOAD=$PWD/sendmsg.so host twitter.com 8.8.8.8
  LD_PRELOAD=$PWD/sendmsg.so dig twitter.com @8.8.8.8
*/
#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>

static ssize_t (*_sendmsg)(int, const struct msghdr *, int);

ssize_t sendmsg(int sockfd, const struct msghdr *msg, int flags)
{
	if (!_sendmsg) {
		_sendmsg = dlsym(RTLD_NEXT, "sendmsg");
		if (!_sendmsg) {
			perror("dlsym");
			exit(1);
		}
	}

	struct msghdr newmsg;
	char buffer[512];
	struct iovec vec;
	if (msg->msg_namelen == sizeof(struct sockaddr_in) && msg->msg_iovlen == 1) {
		const struct sockaddr_in *addr = (const struct sockaddr_in *)msg->msg_name;
		const unsigned char *bytes = msg->msg_iov[0].iov_base;
		const unsigned short *shorts = msg->msg_iov[0].iov_base;
		size_t len = msg->msg_iov[0].iov_len;
		if (addr->sin_family == AF_INET && addr->sin_port == htons(53) &&
		len < sizeof(buffer) - 1 && len > 16 && shorts[2] == htons(1) && shorts[3] == 0) {
			size_t off = 12;
			int ended = 0;
			while (off < len - 4) {
				if (bytes[off] & 0xc0)
					break;
				if (bytes[off] == 0) {
					ended = 1;
					off++;
					break;
				}
				off += 1 + bytes[off];
			}
			if (ended) {
				memcpy(buffer, bytes, off - 1);
				memcpy(buffer + off + 1, bytes + off, len - off);
				buffer[off - 1] = '\xc0';
				buffer[off] = '\x04';

				vec.iov_base = buffer;
				vec.iov_len = len + 1;

				newmsg = *msg;
				newmsg.msg_iov = &vec;
				msg = &newmsg;
			}
		}
	}

	return _sendmsg(sockfd, msg, flags);
}