#include <csignal>
#include <cstdio>
#include <cstdlib>
 
#include <atomic>
#include <thread>
 
#include <errno.h>
#include <error.h>
#include <unistd.h>
 
#include <arpa/inet.h>
#include <netinet/ip.h>
#include <sys/socket.h>
 
/* UDP multicast example. Compile with threads and C++11 (g++ -pthread --std=c++11)
 *
 * Sending to multicast group MCASTIP:MCASTPORT
 *   - create a UDP socket
 *   - send a datagram to MCASTIP:MCASTPORT
 *
 * Receiving datagrams sent to a multicast group MCASTIP:MCASTPORT
 *   - create a UDP socket
 *   - bind to some address and MCASTPORT
 *   - prepare ip_mreqn structure – fill it with MCASTIP and optionally (but recommended) with local IP and interface id
 *   - call setsockopt to set option IP_ADD_MEMBERSHIP on level IPPROTO_IP to the aforementioned structure
 *   - receive datagrams
 *
 * setsockopt of IP_ADD_MEMBERSHIP sends an IGMP message "join group" and tells OS that the program waits for messages for the group.
 *
 * More info at: https://tldp.org/HOWTO/Multicast-HOWTO-6.html
 */
 
const sockaddr_in groupSockaddr{
    .sin_family = AF_INET,
    .sin_port = htons(6789),
    .sin_addr = {.s_addr = inet_addr("239.255.123.45")}
};
 
std::atomic<bool> quitting{false};
 
void ctrl_c(int) {
    if (quitting)
        exit(-1);
    quitting = true;
}
 
void setupSignals();
 
int main() {
 
    // create socket
    int sockfd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd == -1)
        error(1, errno, "socket");
 
    // set SO_REUSEADDR (or only one app will be allowed to receive multicast messages from the groupIP:port pair)
    const int one = 1;
    int res = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
    if (res)
        error(1, errno, "setsockopt SO_REUSEADDR");
 
    // bind
    sockaddr_in myAddress{
        .sin_family = AF_INET,
        .sin_port = groupSockaddr.sin_port,
        .sin_addr = {INADDR_ANY}
    };
    res = bind(sockfd, (sockaddr *)&myAddress, sizeof(myAddress));
    if (res)
        error(1, errno, "bind");
 
    // send IGMP join group
    ip_mreqn groupDescription{
        .imr_multiaddr = groupSockaddr.sin_addr,
        .imr_address = {INADDR_ANY},
        .imr_ifindex = 0
    };
    auto ret = setsockopt(sockfd, IPPROTO_IP, IP_ADD_MEMBERSHIP, &groupDescription, sizeof(groupDescription));
    if (ret)
        error(1, errno, "setsockopt IP_ADD_MEMBERSHIP");
 
    // handle ctrl+c for gracefull exit
    setupSignals();
 
    // receive and display ANY packet that arrives
    std::thread readerThread([&] {
        sockaddr_in peerAddr;
        socklen_t peerAddrSize;
        char buffer[255];
        while (true) {
            peerAddrSize = sizeof(peerAddr);
            auto readSize = recvfrom(sockfd, buffer, 254, MSG_TRUNC, (sockaddr *)&peerAddr, &peerAddrSize);
            if (quitting)
                break;
            if (readSize == -1)
                error(1, errno, "recvfrom");
            if (readSize > 254) {
                error(0, 0, "UDP packet too long - truncated from %ld to %d", readSize, 254);
                readSize = 254;
            }
            buffer[readSize] = 0;
            printf("<%s:%hu> %s", inet_ntoa(peerAddr.sin_addr), ntohs(peerAddr.sin_port), buffer);
        }
    });
 
    // read standard input and send it to the group
    char buffer[255];
    while (true) {
        auto readSize = read(0, buffer, 254);
        if (quitting || !readSize)
            break;
        if (readSize < 0) {
            error(0, errno, "reading stdin");
            break;
        }
 
        auto res = sendto(sockfd, buffer, readSize, 0, (sockaddr *)&groupSockaddr, sizeof(groupSockaddr));
        if (res != readSize) {
            error(0, errno, "sendto");
            break;
        }
    }
 
    quitting = true;
    // leave group
    setsockopt(sockfd, IPPROTO_IP, IP_DROP_MEMBERSHIP, &groupDescription, sizeof(groupDescription));
    // close socket AND interrupt all blocked recvfrom
    shutdown(sockfd, SHUT_RDWR);
    close(sockfd);
    // wait for child
    readerThread.join();
 
    return 0;
}
 
void setupSignals() {
    struct sigaction action;
 
    // read sigint handler description
    auto res = sigaction(SIGINT, nullptr, &action);
    if (res)
        error(1, errno, "sigaction read");
 
    // set new interrupt handler
    action.sa_handler = ctrl_c;
    // disable SA_RESTART flag to exit blocking IO upon a signal
    action.sa_flags &= ~SA_RESTART;
 
    // update sigint handler description
    res = sigaction(SIGINT, &action, nullptr);
    if (res)
        error(1, errno, "sigaction modify");
}