bt.torrent.messaging.MetadataProducer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of bt-core Show documentation
BitTorrent Client Library (Core)
There is a newer version: 1.10
package bt.torrent.messaging;

import bt.magnet.UtMetadata;
import bt.metainfo.Torrent;
import bt.net.Peer;
import bt.protocol.Message;
import bt.runtime.Config;
import bt.torrent.annotation.Consumes;
import bt.torrent.annotation.Produces;

import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.function.Consumer;
import java.util.function.Supplier;

public class MetadataProducer {

    private final Supplier torrentSupplier;

    // initialized on the first metadata request if the torrent is present
    private volatile ExchangedMetadata metadata;

    private final ConcurrentMap> outboundMessages;

    private final int metadataExchangeBlockSize;

    public MetadataProducer(Supplier torrentSupplier,
                            Config config) {
        this.torrentSupplier = torrentSupplier;
        this.outboundMessages = new ConcurrentHashMap<>();
        this.metadataExchangeBlockSize = config.getMetadataExchangeBlockSize();
    }

    @Consumes
    public void consume(UtMetadata message, MessageContext context) {
        Peer peer = context.getPeer();
        // being lenient herer and not checking if the peer advertised ut_metadata support
        switch (message.getType()) {
            case REQUEST: {
                // TODO: spam protection
                processMetadataRequest(peer, message.getPieceIndex());
            }
            default: {
                // ignore
            }
        }
    }

    private void processMetadataRequest(Peer peer, int pieceIndex) {
        Message response;

        Torrent torrent = torrentSupplier.get();
        if (torrent == null) {
            // reject all requests as we don't have the torrent yet
            response = UtMetadata.reject(pieceIndex);
        } else {
            if (metadata == null) {
                metadata = new ExchangedMetadata(torrent.getSource().getExchangedMetadata(), metadataExchangeBlockSize);
            }

            response = UtMetadata.data(pieceIndex, metadata.length(), metadata.getBlock(pieceIndex));
        }

        getOrCreateOutboundMessages(peer).add(response);
    }

    private Queue getOrCreateOutboundMessages(Peer peer) {
        Queue queue = outboundMessages.get(peer);
        if (queue == null) {
            queue = new LinkedBlockingQueue<>();
            Queue existing = outboundMessages.putIfAbsent(peer, queue);
            if (existing != null) {
                queue = existing;
            }
        }
        return queue;
    }

    @Produces
    public void produce(Consumer messageConsumer, MessageContext context) {
        Peer peer = context.getPeer();

        Queue queue = outboundMessages.get(peer);
        if (queue != null && queue.size() > 0) {
            messageConsumer.accept(queue.poll());
        }
    }
}