bt.torrent.messaging.MetadataConsumer 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.IMetadataService;
import bt.metainfo.Torrent;
import bt.metainfo.TorrentId;
import bt.net.Peer;
import bt.protocol.Message;
import bt.protocol.extended.ExtendedHandshake;
import bt.runtime.Config;
import bt.torrent.annotation.Consumes;
import bt.torrent.annotation.Produces;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Arrays;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;

public class MetadataConsumer {
    private static final Logger LOGGER = LoggerFactory.getLogger(MetadataConsumer.class);

    private static final Duration FIRST_BLOCK_ARRIVAL_TIMEOUT = Duration.ofSeconds(10);
    private static final Duration WAIT_BEFORE_REREQUESTING_AFTER_REJECT = Duration.ofSeconds(10);

    private final ConcurrentMap peersWithoutMetadata;

    private final Set supportingPeers;
    private final ConcurrentMap requestedFirstPeers;
    private final Set requestedAllPeers;

    private volatile ExchangedMetadata metadata;

    private final IMetadataService metadataService;

    private final TorrentId torrentId;

    // set immediately after metadata has been fetched and verified
    private final AtomicReference torrent;

    private final int metadataExchangeBlockSize;
    private final int metadataExchangeMaxSize;

    public MetadataConsumer(IMetadataService metadataService,
                            TorrentId torrentId,
                            Config config) {

        this.peersWithoutMetadata = new ConcurrentHashMap<>();

        this.supportingPeers = ConcurrentHashMap.newKeySet();
        this.requestedFirstPeers = new ConcurrentHashMap<>();
        this.requestedAllPeers = ConcurrentHashMap.newKeySet();

        this.metadataService = metadataService;

        this.torrentId = Objects.requireNonNull(torrentId);
        this.torrent = new AtomicReference<>();

        this.metadataExchangeBlockSize = config.getMetadataExchangeBlockSize();
        this.metadataExchangeMaxSize = config.getMetadataExchangeMaxSize();
    }

    @Consumes
    public void consume(ExtendedHandshake handshake, MessageContext messageContext) {
        if (handshake.getSupportedMessageTypes().contains("ut_metadata")) {
            // TODO: peer may eventually turn off the ut_metadata extension
            // moreover the extended handshake message type map is additive,
            // so we can't learn about the peer turning off extensions solely from the message
            supportingPeers.add(messageContext.getPeer());
        }
    }

    @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 DATA: {
                int totalSize = message.getTotalSize().get();
                if (totalSize >= metadataExchangeMaxSize) {
                    throw new IllegalStateException("Declared metadata size is too large: " + totalSize +
                            "; max allowed is " + metadataExchangeMaxSize);
                }
                processMetadataBlock(message.getPieceIndex(), totalSize, message.getData().get());
            }
            case REJECT: {
                peersWithoutMetadata.put(peer, System.currentTimeMillis());
            }
            default: {
                // ignore
            }
        }
    }

    private void processMetadataBlock(int pieceIndex, int totalSize, byte[] data) {
        if (metadata == null) {
            metadata = new ExchangedMetadata(totalSize, metadataExchangeBlockSize);
        }

        if (!metadata.isBlockPresent(pieceIndex)) {
            metadata.setBlock(pieceIndex, data);

            if (metadata.isComplete()) {
                byte[] digest = metadata.getSha1Digest();
                if (Arrays.equals(digest, torrentId.getBytes())) {
                    Torrent fetchedTorrent = null;
                    try {
                        fetchedTorrent = metadataService.fromByteArray(metadata.getBytes());
                    } catch (Exception e) {
                        LOGGER.error("Processing of metadata failed: " + torrentId, e);
                        metadata = null;
                    }

                    if (fetchedTorrent != null) {
                        synchronized (torrent) {
                            torrent.set(fetchedTorrent);
                            requestedFirstPeers.clear();
                            requestedAllPeers.clear();
                            torrent.notifyAll();
                        }
                    }
                } else {
                    LOGGER.warn("Metadata fetched, but hash does not match the torrent ID: {}. Will re-fetch", torrentId);
                    // restart the process
                    // TODO: terminate peer connections that the metadata was fetched from?
                    // or just try again with the others?
                    metadata = null;
                }
            }
        }
    }

    @Produces
    public void produce(Consumer messageConsumer, MessageContext context) {
        // stop here if metadata has already been fetched
        if (torrent.get() != null) {
            return;
        }

        Peer peer = context.getPeer();
        if (supportingPeers.contains(peer)) {
            if (peersWithoutMetadata.containsKey(peer)) {
                if ((System.currentTimeMillis() - peersWithoutMetadata.get(peer)) >= WAIT_BEFORE_REREQUESTING_AFTER_REJECT.toMillis()) {
                    peersWithoutMetadata.remove(peer);
                }
            }

            if (!peersWithoutMetadata.containsKey(peer)) {
                if (metadata == null) {
                    if (!requestedFirstPeers.containsKey(peer) ||
                            (System.currentTimeMillis() - requestedFirstPeers.get(peer) > FIRST_BLOCK_ARRIVAL_TIMEOUT.toMillis())) {
                        requestedFirstPeers.put(peer, System.currentTimeMillis());
                        // start with the first piece of metadata
                        messageConsumer.accept(UtMetadata.request(0));
                    }
                } else if (!requestedAllPeers.contains(peer)) {
                    requestedAllPeers.add(peer);
                    // TODO: larger metadata should be handled in more intelligent way
                    // starting with block #1 because by now we should have already received block #0
                    for (int i = 1; i < metadata.getBlockCount(); i++) {
                        messageConsumer.accept(UtMetadata.request(i));
                    }
                }
            }
        }
    }

    /**
     * @return Torrent, blocking the calling thread if it hasn't been fetched yet
     */
    public Torrent waitForTorrent() {
        if (torrent.get() == null) {
            synchronized (torrent) {
                if (torrent.get() == null) {
                    try {
                        torrent.wait();
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
        return torrent.get();
    }
}