screen_cast/src/network/discovery.cpp

#include "screencast/network/discovery.h"

#include <avahi-client/client.h>
#include <avahi-client/lookup.h>
#include <avahi-client/publish.h>
#include <avahi-common/address.h>
#include <avahi-common/alternative.h>
#include <avahi-common/error.h>
#include <avahi-common/malloc.h>
#include <avahi-common/thread-watch.h>

#include <mutex>
#include <string>
#include <utility>

namespace sc {
namespace {

// DNS-SD service type for screencast receivers. The service advertises the
// signaling (TCP) endpoint; the media endpoint is negotiated in the session.
constexpr const char* kServiceType = "_screencast._tcp";

// Threading model: all Avahi API calls made from application threads run
// under avahi_threaded_poll_lock; callbacks run on the Avahi poll thread and
// only take state_mutex_, never the poll lock, so lock ordering is strict
// (poll lock -> state_mutex_) and circular waits are impossible.
class AvahiDiscovery final : public DiscoveryService {
  public:
    AvahiDiscovery() = default;

    ~AvahiDiscovery() override {
        stop();
    }

    AvahiDiscovery(const AvahiDiscovery&) = delete;
    AvahiDiscovery& operator=(const AvahiDiscovery&) = delete;

    bool start() {
        poll_ = avahi_threaded_poll_new();
        if (poll_ == nullptr) {
            return false;
        }

        int error = 0;
        // The client is created before the poll thread starts, as Avahi
        // recommends; a failure here usually means avahi-daemon is down.
        client_ = avahi_client_new(avahi_threaded_poll_get(poll_),
                                   static_cast<AvahiClientFlags>(0),
                                   &AvahiDiscovery::on_client_state,
                                   this,
                                   &error);
        if (client_ == nullptr) {
            std::lock_guard lock(state_mutex_);
            last_error_ = std::string{"avahi_client_new failed: "} + avahi_strerror(error);
            avahi_threaded_poll_free(poll_);
            poll_ = nullptr;
            return false;
        }

        avahi_threaded_poll_start(poll_);
        return true;
    }

    bool announce(const std::string& service_name, uint16_t signaling_port) override {
        {
            std::lock_guard lock(state_mutex_);
            announced_name_ = service_name;
            announced_port_ = signaling_port;
            wants_announce_ = true;
        }
        // Commit immediately if the client is already up; otherwise the
        // client-state callback commits once Avahi reaches RUNNING and this
        // returns true (deferred). A commit failure surfaces as false.
        avahi_threaded_poll_lock(poll_);
        bool committed = true;
        if (avahi_client_get_state(client_) == AVAHI_CLIENT_S_RUNNING) {
            std::lock_guard lock(state_mutex_);
            committed = create_group_locked();
        }
        avahi_threaded_poll_unlock(poll_);
        return committed;
    }

    bool browse(PeerCallback on_peer) override {
        {
            std::lock_guard lock(state_mutex_);
            peer_callback_ = std::move(on_peer);
        }
        avahi_threaded_poll_lock(poll_);
        browser_ = avahi_service_browser_new(client_,
                                             AVAHI_IF_UNSPEC,
                                             AVAHI_PROTO_UNSPEC,
                                             kServiceType,
                                             nullptr,
                                             static_cast<AvahiLookupFlags>(0),
                                             &AvahiDiscovery::on_browser_event,
                                             this);
        avahi_threaded_poll_unlock(poll_);
        return browser_ != nullptr;
    }

    void stop() override {
        if (browser_ != nullptr) {
            avahi_threaded_poll_lock(poll_);
            avahi_service_browser_free(browser_);
            browser_ = nullptr;
            avahi_threaded_poll_unlock(poll_);
        }
        if (group_ != nullptr) {
            avahi_threaded_poll_lock(poll_);
            avahi_entry_group_free(group_);
            group_ = nullptr;
            avahi_threaded_poll_unlock(poll_);
        }
        if (client_ != nullptr) {
            avahi_threaded_poll_lock(poll_);
            avahi_client_free(client_);
            client_ = nullptr;
            avahi_threaded_poll_unlock(poll_);
        }
        if (poll_ != nullptr) {
            avahi_threaded_poll_stop(poll_);
            avahi_threaded_poll_free(poll_);
            poll_ = nullptr;
        }
    }

    std::string last_error() const override {
        std::lock_guard lock(state_mutex_);
        return last_error_;
    }

  private:
    static void on_client_state(AvahiClient* /*client*/, AvahiClientState state, void* userdata) {
        auto* self = static_cast<AvahiDiscovery*>(userdata);
        switch (state) {
            case AVAHI_CLIENT_S_RUNNING:
                self->on_client_running();
                break;
            case AVAHI_CLIENT_FAILURE: {
                std::lock_guard lock(self->state_mutex_);
                self->last_error_ = "avahi client failure";
                break;
            }
            case AVAHI_CLIENT_S_REGISTERING:
            case AVAHI_CLIENT_S_COLLISION:
            default:
                break; // transient server-wide states
        }
    }

    void on_client_running() {
        std::lock_guard lock(state_mutex_);
        if (wants_announce_ && group_ == nullptr) {
            (void)create_group_locked();
        }
    }

    bool create_group_locked() {
        if (group_ != nullptr) {
            return true;
        }
        group_ = avahi_entry_group_new(client_, &AvahiDiscovery::on_group_state, this);
        if (group_ == nullptr) {
            last_error_ = "avahi_entry_group_new failed";
            return false;
        }

        int result = avahi_entry_group_add_service(group_,
                                                   AVAHI_IF_UNSPEC,
                                                   AVAHI_PROTO_UNSPEC,
                                                   static_cast<AvahiPublishFlags>(0),
                                                   announced_name_.c_str(),
                                                   kServiceType,
                                                   nullptr,
                                                   nullptr,
                                                   announced_port_,
                                                   nullptr);
        if (result == AVAHI_ERR_COLLISION) {
            char* alternative = avahi_alternative_service_name(announced_name_.c_str());
            if (alternative != nullptr) {
                announced_name_ = alternative;
                avahi_free(alternative);
                result = avahi_entry_group_add_service(group_,
                                                       AVAHI_IF_UNSPEC,
                                                       AVAHI_PROTO_UNSPEC,
                                                       static_cast<AvahiPublishFlags>(0),
                                                       announced_name_.c_str(),
                                                       kServiceType,
                                                       nullptr,
                                                       nullptr,
                                                       announced_port_,
                                                       nullptr);
            }
        }
        if (result < 0 || avahi_entry_group_commit(group_) < 0) {
            last_error_ = std::string{"announcing the service failed: "} + avahi_strerror(result);
            avahi_entry_group_reset(group_);
            avahi_entry_group_free(group_);
            group_ = nullptr;
            return false;
        }
        return true;
    }

    static void on_group_state(AvahiEntryGroup* group, AvahiEntryGroupState state, void* userdata) {
        auto* self = static_cast<AvahiDiscovery*>(userdata);
        switch (state) {
            case AVAHI_ENTRY_GROUP_COLLISION: {
                std::lock_guard lock(self->state_mutex_);
                avahi_entry_group_reset(group);
                char* alternative = avahi_alternative_service_name(self->announced_name_.c_str());
                if (alternative != nullptr) {
                    self->announced_name_ = alternative;
                    avahi_free(alternative);
                }
                (void)self->create_group_locked();
                break;
            }
            case AVAHI_ENTRY_GROUP_FAILURE:
                self->set_last_error("service registration failed");
                break;
            default:
                break;
        }
    }

    static void on_browser_event(AvahiServiceBrowser* /*browser*/,
                                 AvahiIfIndex interface,
                                 AvahiProtocol protocol,
                                 AvahiBrowserEvent event,
                                 const char* name,
                                 const char* type,
                                 const char* domain,
                                 AvahiLookupResultFlags /*flags*/,
                                 void* userdata) {
        auto* self = static_cast<AvahiDiscovery*>(userdata);
        switch (event) {
            case AVAHI_BROWSER_NEW:
                // Ownership of the resolver is ours to release in its callback.
                if (avahi_service_resolver_new(self->client_,
                                               interface,
                                               protocol,
                                               name,
                                               type,
                                               domain,
                                               AVAHI_PROTO_UNSPEC,
                                               static_cast<AvahiLookupFlags>(0),
                                               &AvahiDiscovery::on_resolved,
                                               self) == nullptr) {
                    self->set_last_error("creating a service resolver failed");
                }
                break;
            case AVAHI_BROWSER_FAILURE:
                self->set_last_error("service browsing failed");
                break;
            case AVAHI_BROWSER_REMOVE:
                // Removals are not tracked in this MVP; peers accumulate for
                // the duration of the browse window.
            case AVAHI_BROWSER_CACHE_EXHAUSTED:
            case AVAHI_BROWSER_ALL_FOR_NOW:
            default:
                break;
        }
    }

    static void on_resolved(AvahiServiceResolver* resolver,
                            AvahiIfIndex /*interface*/,
                            AvahiProtocol /*protocol*/,
                            AvahiResolverEvent event,
                            const char* name,
                            const char* /*type*/,
                            const char* /*domain*/,
                            const char* /*host_name*/,
                            const AvahiAddress* address,
                            uint16_t port,
                            AvahiStringList* /*txt*/,
                            AvahiLookupResultFlags /*flags*/,
                            void* userdata) {
        auto* self = static_cast<AvahiDiscovery*>(userdata);
        if (event == AVAHI_RESOLVER_FOUND && address != nullptr) {
            char address_text[AVAHI_ADDRESS_STR_MAX];
            avahi_address_snprint(address_text, sizeof(address_text), address);

            DiscoveredPeer peer;
            peer.service_name = name != nullptr ? name : "";
            peer.host = address_text;
            peer.signaling_port = port;

            PeerCallback callback;
            {
                std::lock_guard lock(self->state_mutex_);
                callback = self->peer_callback_;
            }
            if (callback != nullptr) {
                callback(peer);
            }
        }
        avahi_service_resolver_free(resolver);
    }

    void set_last_error(std::string message) {
        std::lock_guard lock(state_mutex_);
        last_error_ = std::move(message);
    }

    AvahiThreadedPoll* poll_ = nullptr;
    AvahiClient* client_ = nullptr;
    AvahiEntryGroup* group_ = nullptr;
    AvahiServiceBrowser* browser_ = nullptr;

    mutable std::mutex state_mutex_;
    std::string last_error_;
    std::string announced_name_;
    uint16_t announced_port_ = 0;
    bool wants_announce_ = false;
    PeerCallback peer_callback_;
};

} // namespace

NetworkResult<std::unique_ptr<DiscoveryService>> DiscoveryFactory::create_avahi() {
    auto discovery = std::make_unique<AvahiDiscovery>();
    if (!discovery->start()) {
        return NetworkError{discovery->last_error()};
    }
    return std::unique_ptr<DiscoveryService>(std::move(discovery));
}

} // namespace sc