/*
 * Send and receive multicast packets.
 *
 * The first send/receive pair should succeed.
 * The second should block on the receive since the source is blocked.
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <strings.h>
#include <string.h>
#include <unistd.h>

#define EXAMPLE_PORT 6000
#define EXAMPLE_GROUP "239.0.0.1"

#define MULTICAST_GROUP "225.4.5.6"
#define MULTICAST_PORT 13534

static
int sendDatagram(char* fromAddress,
                 char* groupAddress,
                 short port) {

    /* set up socket */
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock == -1) {
        perror("socket");
        return -1;
    }

    /* Group */
    struct sockaddr_in group;
    memset((char *) &group, 0, sizeof(group));
    group.sin_family = AF_INET;
    group.sin_addr.s_addr = inet_addr(groupAddress);
    group.sin_port = htons(port);

    /* Local */
    struct in_addr localInterface;
    localInterface.s_addr = inet_addr(fromAddress);
    if (setsockopt(sock, IPPROTO_IP, IP_MULTICAST_IF,
                   (char *)&localInterface, sizeof(localInterface)) == -1) {
        perror("Local");
        return -1;
    }

    /* Send */
    int message = rand();
    if (sendto(sock, &message, sizeof(int), 0,
               (struct sockaddr*) &group, sizeof(group)) == -1) {
        perror("Send");
        return -1;
    }

    return message;
}

int
createReceiveSocket(char* localAddress,
                    char* groupAddress,
                    short port) {

    /* Socket */
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock == -1) {
        perror("socket");
        return -1;
    }

    /* Enable reuse */
    int reuse = 1;
    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR,
                   (char *)&reuse, sizeof(reuse)) == -1) {
        perror("Reuse");
        close(sock);
        return -1;
    }

    /* Bind */
    struct sockaddr_in local;
    memset((char *) &local, 0, sizeof(local));
    local.sin_family = AF_INET;
    local.sin_port = htons(port);
    local.sin_addr.s_addr = INADDR_ANY;
    if (bind(sock, (struct sockaddr*)&local, sizeof(local)) == -1) {
        perror("Bind");
        close(sock);
        return -1;
    }

    /* Join group */
    struct ip_mreq group;
    group.imr_multiaddr.s_addr = inet_addr(groupAddress);
    group.imr_interface.s_addr = inet_addr(localAddress);
    if (setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP,
                   (char *)&group, sizeof(group)) == -1) {
        perror("Join");
        close(sock);
        return -1;
    }

    return sock;
}

int
receiveDatagram(int sock,
                int expectedMessage) {
    struct sockaddr_in source;
    socklen_t sourcelen = sizeof(source);

    int message;
    int count = recvfrom(sock, &message, sizeof(message), 0,
                         (struct sockaddr *) &source, &sourcelen);
    if (count == -1) {
       perror("Receive");
       return -1;
    }

    printf("Received count: %d\n", count);
    printf("Source: %s\n", inet_ntoa(source.sin_addr));
    printf("Got message: %d, Expected message: %d\n", message, expectedMessage);

    return count;
}

void usage() {
    fprintf(stderr, "Usage: multicast <multicast capable local interface address>\n");
    fprintf(stderr, "E.g. multicast 172.16.0.1\n");
}

int main(int argc, char* argv[]) {

    if (argc < 2) {
        usage();
        exit(1);
    }
    printf("Using interface %s\n", argv[1]);

    /* Create socket */
    int sock = createReceiveSocket(argv[1], MULTICAST_GROUP, MULTICAST_PORT);
    printf("Created receive socket %d\n", sock);

    /* Send/receive */
    int messageId = sendDatagram(argv[1], MULTICAST_GROUP, MULTICAST_PORT);
    printf("Sent message %d\n", messageId);
    receiveDatagram(sock, messageId);

    /* Squelch */
    struct ip_mreq_source mreqs;
    mreqs.imr_multiaddr.s_addr = inet_addr(MULTICAST_GROUP);
    mreqs.imr_sourceaddr.s_addr = inet_addr(argv[1]);
    mreqs.imr_interface.s_addr = inet_addr(argv[1]);
    printf("Blocking source %s\n", argv[1]);
    if (setsockopt(sock, IPPROTO_IP, IP_BLOCK_SOURCE, &mreqs, sizeof(mreqs)) == -1) {
        perror("Blocking source");
	exit(1);
    }

    /* Send/receive */
    messageId = sendDatagram(argv[1], MULTICAST_GROUP, MULTICAST_PORT);
    printf("Sent message %d\n", messageId);
    printf("Program should hang here since the source is blocked.\n");
    receiveDatagram(sock, messageId);
}