diff --git a/btpd/net.c b/btpd/net.c index e627378..719c57c 100644 --- a/btpd/net.c +++ b/btpd/net.c @@ -476,12 +476,7 @@ read_bitfield(struct peer *p, unsigned long rmax) rd->iob.buf_off += nread; if (rd->iob.buf_off == rd->iob.buf_len) { - bcopy(rd->iob.buf, p->piece_field, rd->iob.buf_len); - for (unsigned i = 0; i < p->tp->meta.npieces; i++) - if (has_bit(p->piece_field, i)) { - p->npieces++; - cm_on_piece_ann(p, i); - } + peer_on_bitfield(p, rd->iob.buf); free(rd); net_generic_reader(p); } else @@ -513,16 +508,7 @@ read_piece(struct peer *p, unsigned long rmax) p->rate_to_me[btpd.seconds % RATEHISTORY] += nread; p->tp->downloaded += nread; if (rd->iob.buf_off == rd->iob.buf_len) { - struct piece_req *req = BTPDQ_FIRST(&p->my_reqs); - if (req != NULL && - req->index == rd->index && - req->begin == rd->begin && - req->length == rd->iob.buf_len) { - // - off_t cbegin = rd->index * p->tp->meta.piece_length + rd->begin; - torrent_put_bytes(p->tp, rd->iob.buf, cbegin, rd->iob.buf_len); - cm_on_block(p); - } + peer_on_piece(p, rd->index, rd->begin, rd->iob.buf_len, rd->iob.buf); free(rd); net_generic_reader(p); } else @@ -578,54 +564,33 @@ net_generic_read(struct peer *p, unsigned long rmax) btpd_log(BTPD_L_MSG, "choke.\n"); if (msg_len != 1) goto bad_data; - if ((p->flags & (PF_P_CHOKE|PF_I_WANT)) == PF_I_WANT) { - p->flags |= PF_P_CHOKE; - cm_on_undownload(p); - } else - p->flags |= PF_P_CHOKE; + peer_on_choke(p); break; case MSG_UNCHOKE: btpd_log(BTPD_L_MSG, "unchoke.\n"); if (msg_len != 1) goto bad_data; - if ((p->flags & (PF_P_CHOKE|PF_I_WANT)) - == (PF_P_CHOKE|PF_I_WANT)) { - p->flags &= ~PF_P_CHOKE; - cm_on_download(p); - } else - p->flags &= ~PF_P_CHOKE; + peer_on_unchoke(p); break; case MSG_INTEREST: btpd_log(BTPD_L_MSG, "interested.\n"); if (msg_len != 1) goto bad_data; - if ((p->flags & (PF_P_WANT|PF_I_CHOKE)) == 0) { - p->flags |= PF_P_WANT; - cm_on_upload(p); - } else - p->flags |= PF_P_WANT; + peer_on_interest(p); break; case MSG_UNINTEREST: btpd_log(BTPD_L_MSG, "uninterested.\n"); if (msg_len != 1) goto bad_data; - if ((p->flags & (PF_P_WANT|PF_I_CHOKE)) == PF_P_WANT) { - p->flags &= ~PF_P_WANT; - cm_on_unupload(p); - } else - p->flags &= ~PF_P_WANT; + peer_on_uninterest(p); break; case MSG_HAVE: btpd_log(BTPD_L_MSG, "have.\n"); if (msg_len != 5) goto bad_data; else if (len - off >= msg_len + 4) { - unsigned long piece = net_read32(buf + off + 5); - if (!has_bit(p->piece_field, piece)) { - set_bit(p->piece_field, piece); - p->npieces++; - cm_on_piece_ann(p, piece); - } + uint32_t index = net_read32(buf + off + 5); + peer_on_have(p, index); } else got_part = 1; break; @@ -635,14 +600,9 @@ net_generic_read(struct peer *p, unsigned long rmax) goto bad_data; else if (p->npieces != 0) goto bad_data; - else if (len - off >= msg_len + 4) { - bcopy(buf + off + 5, p->piece_field, msg_len - 1); - for (unsigned i = 0; i < p->tp->meta.npieces; i++) - if (has_bit(p->piece_field, i)) { - p->npieces++; - cm_on_piece_ann(p, i); - } - } else { + else if (len - off >= msg_len + 4) + peer_on_bitfield(p, buf + off + 5); + else { struct bitfield_reader *rp; size_t mem = sizeof(*rp) + msg_len - 1; p->reader->kill(p->reader); @@ -666,23 +626,18 @@ net_generic_read(struct peer *p, unsigned long rmax) if ((p->flags & (PF_P_WANT|PF_I_CHOKE)) != PF_P_WANT) break; uint32_t index, begin, length; - off_t cbegin; - char *content; index = net_read32(buf + off + 5); begin = net_read32(buf + off + 9); length = net_read32(buf + off + 13); if (length > (1 << 15)) goto bad_data; - if (index >= p->tp->meta.npieces - || !has_bit(p->tp->piece_field, index)) + if (index >= p->tp->meta.npieces) goto bad_data; - if (begin + length > p->tp->meta.piece_length) + if (!has_bit(p->tp->piece_field, index)) goto bad_data; - cbegin = index * p->tp->meta.piece_length + begin; - if (cbegin + length > p->tp->meta.total_length) + if (begin + length > torrent_piece_size(p->tp, index)) goto bad_data; - content = torrent_get_bytes(p->tp, cbegin, length); - net_send_piece(p, index, begin, content, length); + peer_on_request(p, index, begin, length); } else got_part = 1; break; @@ -695,18 +650,9 @@ net_generic_read(struct peer *p, unsigned long rmax) uint32_t begin = net_read32(buf + off + 9); uint32_t length = msg_len - 9; if (len - off >= msg_len + 4) { - off_t cbegin = index * p->tp->meta.piece_length + begin; p->tp->downloaded += length; p->rate_to_me[btpd.seconds % RATEHISTORY] += length; - struct piece_req *req = BTPDQ_FIRST(&p->my_reqs); - if (req != NULL && - req->index == index && - req->begin == begin && - req->length == length) { - // - torrent_put_bytes(p->tp, buf + off + 13, cbegin, length); - cm_on_block(p); - } + peer_on_piece(p, index, begin, length, buf + off + 13); } else { struct piece_reader *rp; size_t mem = sizeof(*rp) + length; @@ -734,27 +680,16 @@ net_generic_read(struct peer *p, unsigned long rmax) if (msg_len != 13) goto bad_data; else if (len - off >= msg_len + 4) { - struct piece_req *req; - uint32_t index, begin, length; - - index = net_read32(buf + off + 5); - begin = net_read32(buf + off + 9); - length = net_read32(buf + off + 13); - + uint32_t index = net_read32(buf + off + 5); + uint32_t begin = net_read32(buf + off + 9); + uint32_t length = net_read32(buf + off + 13); + if (index > p->tp->meta.npieces) + goto bad_data; + if (begin + length > torrent_piece_size(p->tp, index)) + goto bad_data; btpd_log(BTPD_L_MSG, "cancel: %u, %u, %u\n", index, begin, length); - - req = BTPDQ_FIRST(&p->p_reqs); - while (req != NULL) { - if (req->index == index && - req->begin == begin && - req->length == length) { - btpd_log(BTPD_L_MSG, "cancel matched.\n"); - net_unsend_piece(p, req); - break; - } - req = BTPDQ_NEXT(req, entry); - } + peer_on_cancel(p, index, begin, length); } else got_part = 1; break; diff --git a/btpd/net.h b/btpd/net.h index 7216362..bc27999 100644 --- a/btpd/net.h +++ b/btpd/net.h @@ -93,6 +93,8 @@ void net_send_choke(struct peer *p); void net_send_have(struct peer *p, uint32_t index); void net_send_request(struct peer *p, struct piece_req *req); +void net_send_piece(struct peer *p, uint32_t index, uint32_t begin, + char *block, size_t blen); void net_send_cancel(struct peer *p, struct piece_req *req); void net_handshake(struct peer *p, int incoming); diff --git a/btpd/peer.c b/btpd/peer.c index 9c2dc90..b8646d7 100644 --- a/btpd/peer.c +++ b/btpd/peer.c @@ -1,6 +1,8 @@ #include #include #include + +#include #include #include @@ -166,10 +168,8 @@ peer_create_in(int sd) } void -peer_create_out(struct torrent *tp, - const uint8_t *id, - const char *ip, - int port) +peer_create_out(struct torrent *tp, const uint8_t *id, + const char *ip, int port) { int sd; struct peer *p; @@ -201,3 +201,106 @@ peer_create_out_compact(struct torrent *tp, const char *compact) p->tp = tp; net_handshake(p, 0); } + +void +peer_on_choke(struct peer *p) +{ + if ((p->flags & (PF_P_CHOKE|PF_I_WANT)) == PF_I_WANT) { + p->flags |= PF_P_CHOKE; + cm_on_undownload(p); + } else + p->flags |= PF_P_CHOKE; +} + +void +peer_on_unchoke(struct peer *p) +{ + if ((p->flags & (PF_P_CHOKE|PF_I_WANT)) == (PF_P_CHOKE|PF_I_WANT)) { + p->flags &= ~PF_P_CHOKE; + cm_on_download(p); + } else + p->flags &= ~PF_P_CHOKE; +} + +void +peer_on_interest(struct peer *p) +{ + if ((p->flags & (PF_P_WANT|PF_I_CHOKE)) == 0) { + p->flags |= PF_P_WANT; + cm_on_upload(p); + } else + p->flags |= PF_P_WANT; +} + +void +peer_on_uninterest(struct peer *p) +{ + if ((p->flags & (PF_P_WANT|PF_I_CHOKE)) == PF_P_WANT) { + p->flags &= ~PF_P_WANT; + cm_on_unupload(p); + } else + p->flags &= ~PF_P_WANT; +} + +void +peer_on_have(struct peer *p, uint32_t index) +{ + if (!has_bit(p->piece_field, index)) { + set_bit(p->piece_field, index); + p->npieces++; + cm_on_piece_ann(p, index); + } +} + +void +peer_on_bitfield(struct peer *p, uint8_t *field) +{ + assert(p->npieces == 0); + bcopy(field, p->piece_field, (size_t)ceil(p->tp->meta.npieces / 8.0)); + for (uint32_t i = 0; i < p->tp->meta.npieces; i++) { + if (has_bit(p->piece_field, i)) { + p->npieces++; + cm_on_piece_ann(p, i); + } + } +} + +void +peer_on_piece(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length, const char *data) +{ + off_t cbegin = index * p->tp->meta.piece_length + begin; + struct piece_req *req = BTPDQ_FIRST(&p->my_reqs); + if (req != NULL && + req->index == index && + req->begin == begin && + req->length == length) { + torrent_put_bytes(p->tp, data, cbegin, length); + cm_on_block(p); + } +} + +void +peer_on_request(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length) +{ + off_t cbegin = index * p->tp->meta.piece_length + begin; + char * content = torrent_get_bytes(p->tp, cbegin, length); + net_send_piece(p, index, begin, content, length); +} + +void +peer_on_cancel(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length) +{ + struct piece_req *req = BTPDQ_FIRST(&p->p_reqs); + while (req != NULL) { + if (req->index == index + && req->begin == begin && req->length == length) { + btpd_log(BTPD_L_MSG, "cancel matched.\n"); + net_unsend_piece(p, req); + break; + } + req = BTPDQ_NEXT(req, entry); + } +} diff --git a/btpd/peer.h b/btpd/peer.h index c1e2f75..4547d2a 100644 --- a/btpd/peer.h +++ b/btpd/peer.h @@ -48,7 +48,7 @@ void peer_choke(struct peer *p); void peer_unwant(struct peer *p, uint32_t index); void peer_want(struct peer *p, uint32_t index); void peer_request(struct peer *p, uint32_t index, - uint32_t begin, uint32_t len); + uint32_t begin, uint32_t len); void peer_cancel(struct peer *p, uint32_t index, uint32_t begin, uint32_t len); void peer_have(struct peer *p, uint32_t index); @@ -57,8 +57,21 @@ unsigned long peer_get_rate(unsigned long *rates); void peer_create_in(int sd); void peer_create_out(struct torrent *tp, const uint8_t *id, - const char *ip, int port); + const char *ip, int port); void peer_create_out_compact(struct torrent *tp, const char *compact); void peer_kill(struct peer *p); +void peer_on_interest(struct peer *p); +void peer_on_uninterest(struct peer *p); +void peer_on_choke(struct peer *p); +void peer_on_unchoke(struct peer *p); +void peer_on_have(struct peer *p, uint32_t index); +void peer_on_bitfield(struct peer *p, uint8_t *field); +void peer_on_piece(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length, const char *data); +void peer_on_request(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length); +void peer_on_cancel(struct peer *p, uint32_t index, uint32_t begin, + uint32_t length); + #endif diff --git a/btpd/torrent.c b/btpd/torrent.c index 5419b59..90a463b 100644 --- a/btpd/torrent.c +++ b/btpd/torrent.c @@ -43,7 +43,7 @@ torrent_load3(const char *file, struct metainfo *mi, char *mem, size_t memsiz) btpd_err("Out of memory.\n"); tp->piece_count = btpd_calloc(mi->npieces, sizeof(tp->piece_count[0])); - + BTPDQ_INIT(&tp->peers); BTPDQ_INIT(&tp->getlst); @@ -251,3 +251,14 @@ torrent_get_by_hash(const uint8_t *hash) tp = BTPDQ_NEXT(tp, entry); return tp; } + +off_t +torrent_piece_size(struct torrent *tp, uint32_t index) +{ + if (index < tp->meta.npieces - 1) + return tp->meta.piece_length; + else { + off_t allbutlast = tp->meta.piece_length * (tp->meta.npieces - 1); + return tp->meta.total_length - allbutlast; + } +} diff --git a/btpd/torrent.h b/btpd/torrent.h index c9d5d2c..022201b 100644 --- a/btpd/torrent.h +++ b/btpd/torrent.h @@ -64,4 +64,6 @@ int torrent_has_peer(struct torrent *tp, const uint8_t *id); struct torrent *torrent_get_by_hash(const uint8_t *hash); +off_t torrent_piece_size(struct torrent *tp, uint32_t index); + #endif