messagesQueue = new LinkedList<>();
28 | private final DhtServerConfig config;
29 | private final Random random = new SecureRandom();
30 | private final byte[] token = new byte[2];
31 |
32 | @Autowired
33 | public DhtMainChain(DhtServerConfig config) {
34 | this.config = config;
35 | }
36 |
37 | private void addMessageQueue(KrpcMessage message) {
38 | messagesQueue.add(message);
39 | }
40 |
41 | /**
42 | * Maximum weight
43 | *
44 | * @return 10
45 | */
46 | @Override
47 | public int weights() {
48 | return 10;
49 | }
50 |
51 | @Override
52 | public KrpcMessage getMessage() throws NullPointerException {
53 | return messagesQueue.remove();
54 | }
55 |
56 | @Override
57 | public boolean isWriteAble() {
58 | return !messagesQueue.isEmpty();
59 | }
60 |
61 | /**
62 | * 基本回复
63 | * {y=r, t=aa, r={id=self_id}}
64 | *
65 | * @param src 请求消息
66 | * @return 基本回复
67 | */
68 | private KrpcMessage baseResponse(KrpcMessage src) {
69 | BencodeMap result = new BencodeMap();
70 | BencodeMap content = new BencodeMap();
71 | content.putByteArray(KrpcToken.ID, config.getSelfNodeId());
72 | result.put(KrpcToken.RESPONSES_MAP, content);
73 | result.put(KrpcToken.TRANSACTION, src.getMessage().get(KrpcToken.TRANSACTION));
74 | result.putString(KrpcToken.TYPE, KrpcToken.TYPE_RESPONSE);
75 | return new KrpcMessage(result, src.getAddress());
76 | }
77 |
78 | @Override
79 | public void onPing(KrpcMessage query) {
80 | addMessageQueue(baseResponse(query));
81 | }
82 |
83 | @Override
84 | public void onFindNodes(KrpcMessage query) {
85 | KrpcMessage response = baseResponse(query);
86 | response.getMessage()
87 | .getBencodeMap(KrpcToken.RESPONSES_MAP)
88 | .put(KrpcToken.NODES, BencodeByteArray.empty());
89 | addMessageQueue(response);
90 | }
91 |
92 | private byte[] nextToken() {
93 | random.nextBytes(token);
94 | return token;
95 | }
96 |
97 | @Override
98 | public void onGetPeer(KrpcMessage query) {
99 | KrpcMessage response = baseResponse(query);
100 | BencodeMap responseMap = response.getMessage().getBencodeMap(KrpcToken.RESPONSES_MAP);
101 | responseMap.put(KrpcToken.NODES, BencodeByteArray.empty());
102 | responseMap.putByteArray(KrpcToken.TOKEN, nextToken());
103 | addMessageQueue(response);
104 | }
105 |
106 | @Override
107 | public void onAnnouncePeer(KrpcMessage query) {
108 | addMessageQueue(baseResponse(query));
109 | }
110 | }
111 |
--------------------------------------------------------------------------------
/server/src/main/java/cc/aguesuka/btfind/dht/KrpcToken.java:
--------------------------------------------------------------------------------
1 | package cc.aguesuka.btfind.dht;
2 |
3 | /**
4 | * http://www.bittorrent.org/beps/bep_0005.html#krpc-protocol
5 | *
6 | * The KRPC protocol is a simple RPC mechanism consisting of bencoded dictionaries sent over UDP.
7 | * UdpHandler single query packet is sent out and a single packet is sent in response.
8 | * There is no retry. There are three message types: query, response, and error.
9 | * For the DHT protocol, there are four queries: ping, find_node, get_peers, and announce_peer.
10 | *
11 | *
12 | * KRPC 消息的类型是 BencodeMap; key 用 string 表示 value 用 String 表示.
13 | *
14 | * @author :aguesuka
15 | * 2019/8/29 15:06
16 | */
17 | @SuppressWarnings("unused")
18 | public class KrpcToken {
19 | public static final int ID_LENGTH = 20;
20 | /**
21 | * Every message has a key "t" with a string value representing a transaction ID.
22 | * This transaction ID is generated by the querying node and is echoed in the response,
23 | * so responses may be correlated with multiple queries to the same node.
24 | * The transaction ID should be encoded as a short string of binary numbers,
25 | * typically 2 characters are enough as they cover 2^16 outstanding queries.
26 | */
27 | public static final String TRANSACTION = "t";
28 | /**
29 | * Every message also has a key "y" with a single character value describing the type of message.
30 | */
31 | public static final String TYPE = "y";
32 | /**
33 | * The value of the "y" key is one of "q" for query
34 | */
35 | public static final String TYPE_QUERY = "q";
36 | /**
37 | * The value of the "y" key is one of "r" for response
38 | */
39 | public static final String TYPE_RESPONSE = "r";
40 | /**
41 | * The value of the "y" key is one of "e" for error
42 | */
43 | public static final String TYPE_ERROR = "e";
44 | /**
45 | * Queries, or KRPC message dictionaries with a "y" value of "q",
46 | * contain two additional keys; "q" and "a".
47 | * Key "a" has a dictionary value containing named arguments to the query.
48 | */
49 | public static final String ARGUMENTS_MAP = "a";
50 | /**
51 | * Key "q" has a string value containing the method name of the query.
52 | */
53 | public static final String QUERY = "q";
54 |
55 | /**
56 | * Responses, or KRPC message dictionaries with a "y" value of "r", contain one additional key "r".
57 | * The value of "r" is a dictionary containing named return values.
58 | * Response messages are sent upon successful completion of a query.
59 | */
60 | public static final String RESPONSES_MAP = "r";
61 |
62 | public static final String PING = "ping";
63 | public static final String FIND_NODE = "find_node";
64 | public static final String GET_PEERS = "get_peers";
65 | public static final String ANNOUNCE_PEER = "announce_peer";
66 | public static final String SAMPLE_INFOHASHES = "sample_infohashes";
67 | public static final String ID = "id";
68 | public static final String TARGET = "target";
69 | public static final String NODES = "nodes";
70 | public static final String INFO_HASH = "info_hash";
71 | public static final String VALUES = "values";
72 | public static final String TOKEN = "token";
73 | public static final String IMPLIED_PORT = "implied_port";
74 | public static final String PORT = "port";
75 |
76 | /**
77 | * There is an optional argument called implied_port which value is either 0 or 1.
78 | * If it is present and non-zero, the port argument should be ignored and the source port of the
79 | * UDP packet should be used as the peer's port instead.
80 | * This is useful for peers behind a NAT that may not know their external port, and supporting uTP, they accept incoming
81 | */
82 | public static final int NOT_IGNORED_PORT = 0;
83 |
84 | }
85 |
--------------------------------------------------------------------------------
/server/pom.xml:
--------------------------------------------------------------------------------
1 |
2 |
5 |
6 | ague-dht
7 | cc.aguesuka
8 | 1.2
9 |
10 | 4.0.0
11 | jar
12 | server
13 |
14 | 2.1.6.RELEASE
15 |
16 |
17 |
18 | cc.aguesuka
19 | bencode
20 | 1.2
21 |
22 |
23 | org.springframework.boot
24 | spring-boot
25 | ${sprint-boot.version}
26 |
27 |
28 | org.springframework.boot
29 | spring-boot-starter
30 | ${sprint-boot.version}
31 |
32 |
33 | org.springframework.boot
34 | spring-boot-configuration-processor
35 | ${sprint-boot.version}
36 | true
37 |
38 |
39 | org.springframework.boot
40 | spring-boot-starter-jdbc
41 | ${sprint-boot.version}
42 |
43 |
44 | org.xerial
45 | sqlite-jdbc
46 | 3.14.2.1
47 |
48 |
49 | org.springframework.boot
50 | spring-boot-starter-test
51 | ${sprint-boot.version}
52 | test
53 |
54 |
55 | org.springframework
56 | spring-test
57 | 5.1.8.RELEASE
58 | test
59 |
60 |
61 |
62 |
63 |
64 | org.springframework.boot
65 | spring-boot-maven-plugin
66 | ${sprint-boot.version}
67 |
68 | cc.aguesuka.btfind.ServiceApplication
69 |
70 |
71 | org.projectlombok
72 | lombok
73 |
74 |
75 | org.springframework.boot
76 | spring-boot-configuration-processor
77 |
78 |
79 | junit
80 | junit
81 |
82 |
83 |
84 |
85 |
86 |
87 | repackage
88 |
89 |
90 |
91 |
92 |
93 |
94 |
95 | lib
96 | BOOT-INF/lib/
97 |
98 | **/*.jar
99 |
100 |
101 |
102 | true
103 | src/main/resources
104 |
105 | *.yml
106 |
107 |
108 |
109 |
110 |
--------------------------------------------------------------------------------
/bencode/src/main/java/cc/aguesuka/bencode/Bencode.java:
--------------------------------------------------------------------------------
1 | package cc.aguesuka.bencode;
2 |
3 | import java.nio.ByteBuffer;
4 | import java.nio.charset.Charset;
5 | import java.nio.charset.StandardCharsets;
6 |
7 | /**
8 | * http://www.bittorrent.org/beps/bep_0003.html#bencoding
9 | *
10 | * bencode是bit-torrent协议的常用编码,消息和种子文件中都有用到.
11 | * 本类提供IBencode对象和byte数组/ByteBuffer的互相转换
12 | * {@link Bencode#toBytes(Object)}
13 | * {@link Bencode#writeToBuffer(IBencode, ByteBuffer)}
14 | * {@link Bencode#parse(ByteBuffer)}
15 | * {@link Bencode#parse(byte[])}
16 | *
17 | * bencode 有四种类型
18 | *
19 | * INTEGER 整数 {@link BencodeInteger}
20 | * Integers are represented by an 'i' followed by the number in base 10 followed by an 'e'.
21 | * For example i3e corresponds to 3 and i-3e corresponds to -3. Integers have no size limitation.
22 | * i-0e is invalid. All encodings with a leading zero, such as i03e,
23 | * are invalid, other than i0e, which of course corresponds to 0.
24 | *
25 | * STRING 对应Java中的byte[]而非字符串; {@link BencodeByteArray} 字符串可以转为BencodeByteArray反之则不一定
26 | * Strings are length-prefixed base ten followed by a colon and the string. For example 4:spam corresponds to 'spam'.
27 | *
28 | * Lists are encoded as an 'l' followed by their elements (also bencoded) followed by an 'e'.
29 | * For example l4:spam4:eggse corresponds to ['spam', 'eggs'].
30 | * LIST 数组 {@link BencodeList}
31 | *
32 | * Dictionaries are encoded as a 'd' followed by a list of alternating keys and their corresponding values followed by an 'e'.
33 | * For example, d3:cow3:moo4:spam4:eggse corresponds to {'cow': 'moo', 'spam': 'eggs'}
34 | * and d4:spaml1:a1:bee corresponds to {'spam': ['a', 'b']}.
35 | * Keys must be strings and appear in sorted order (sorted as raw strings, not alphanumerics).
36 | * DICT 字典 {@link BencodeMap} map的key是字符串
37 | *
38 | * bencode 没有定义null,但是在处理时不做校验,所以解析成对象时,map的最后一个value可能为null,编码为bencode时,null视为空字符串
39 | *
40 | * @author :aguesuka
41 | * 2019/6/26 17:47
42 | */
43 | public final class Bencode {
44 |
45 | private static final Charset charset = StandardCharsets.UTF_8;
46 |
47 | static Charset getCharset() {
48 | return charset;
49 | }
50 |
51 | /**
52 | * 将 ByteBuff 对象转 BencodeMap ,只转换第一个 BencodeMap 并将 ByteBuff 的 offset 置为第一个map的结尾.
53 | *
54 | * @param buff byte buff 对象
55 | * @return BencodeMap
56 | */
57 | public static BencodeMap parse(ByteBuffer buff) {
58 | try {
59 | return new BencodeParser(buff).parser();
60 | } catch (BencodeException e) {
61 | throw e;
62 | } catch (RuntimeException e) {
63 | throw new BencodeException(null, e);
64 | }
65 | }
66 |
67 | /**
68 | * 将 ByteBuff 对象转 BencodeMap ,只转换第一个 BencodeMap 并将 ByteBuff 的 offset 置为第一个map的结尾.
69 | *
70 | * @param bytes byte[] 对象
71 | * @return BencodeMap
72 | */
73 | @SuppressWarnings("WeakerAccess")
74 | public static BencodeMap parse(byte[] bytes) {
75 | try {
76 | return new BencodeParser() {
77 | int position = 0;
78 |
79 | @Override
80 | BencodeException ex(String msg, Throwable ex) {
81 | String m = "at position " + position + ":" + msg;
82 | return new BencodeException(rootMap, m, ex);
83 | }
84 |
85 | @Override
86 | byte byteFromBuffer() {
87 | return bytes[position++];
88 | }
89 |
90 | @Override
91 | byte[] byteArrayFromBuffer(int length) {
92 | byte[] result = new byte[length];
93 | System.arraycopy(bytes, position, result, 0, length);
94 | position += length;
95 | return result;
96 | }
97 | }.parser();
98 | } catch (BencodeException e) {
99 | throw e;
100 | } catch (RuntimeException e) {
101 | throw new BencodeException(null, e);
102 | }
103 |
104 | }
105 |
106 | /**
107 | * 将IBencode对象转为byte数组
108 | *
109 | * @param o IBencode对象
110 | * @return byte 数组
111 | */
112 | static byte[] toBytes(Object o) {
113 | BencodeEncoder bencodeEncode = new BencodeEncoder();
114 | bencodeEncode.putBencode(o);
115 | return bencodeEncode.getResult();
116 | }
117 |
118 | static ByteBuffer writeToBuffer(IBencode o, ByteBuffer buffer) {
119 | BencodeEncoder bencodeEncode = new BencodeEncoder() {
120 | @Override
121 | void writeByteArray(byte[] bytes) {
122 | buffer.put(bytes);
123 | }
124 |
125 | @Override
126 | void writeByte(byte b) {
127 | buffer.put(b);
128 | }
129 |
130 | @Override
131 | byte[] getResult() {
132 | return null;
133 | }
134 | };
135 | bencodeEncode.putBencode(o);
136 | return buffer;
137 | }
138 |
139 | }
140 |
--------------------------------------------------------------------------------
/bencode/src/main/java/cc/aguesuka/bencode/BencodeParser.java:
--------------------------------------------------------------------------------
1 | package cc.aguesuka.bencode;
2 |
3 | import java.nio.ByteBuffer;
4 | import java.util.ArrayDeque;
5 | import java.util.Deque;
6 | import java.util.Objects;
7 |
8 | /**
9 | * @author :aguesuka
10 | * 2019/9/3 20:09
11 | */
12 | class BencodeParser {
13 | BencodeMap rootMap = new BencodeMap();
14 | private ByteBuffer buff;
15 | private boolean isReadMap = true;
16 | private final Deque> stack = new ArrayDeque<>();
17 |
18 | BencodeParser() {
19 | }
20 |
21 | BencodeParser(ByteBuffer buff) {
22 | Objects.requireNonNull(buff);
23 | this.buff = buff;
24 | }
25 |
26 | BencodeMap parser() {
27 | try {
28 |
29 |
30 | stack.add(rootMap);
31 | if (byteFromBuffer() != BencodeToken.DICT) {
32 | throw ex("is not a bencode dict", null);
33 | }
34 | while (!stack.isEmpty()) {
35 | Object last = stack.getLast();
36 | if (isReadMap) {
37 | String key = readStringOrEnd();
38 | if (key != null) {
39 | IBencode bencodeObject = readObject();
40 | Objects.requireNonNull(bencodeObject);
41 | ((BencodeMap) last).put(key, bencodeObject);
42 | }
43 | } else {
44 | IBencode bencode = readObject();
45 | if (bencode != null) {
46 | ((BencodeList) last).add(bencode);
47 | }
48 | }
49 | }
50 | return rootMap;
51 | } catch (BencodeException e) {
52 | throw e;
53 | } catch (RuntimeException e) {
54 | throw ex(e.getMessage(), e);
55 | }
56 | }
57 |
58 | private void onEnd() {
59 | stack.removeLast();
60 | isReadMap = stack.peekLast() instanceof BencodeMap;
61 | }
62 |
63 | private IBencode readObject() {
64 | byte firstByte = byteFromBuffer();
65 | switch (firstByte) {
66 | case BencodeToken.END:
67 | onEnd();
68 | return null;
69 | case BencodeToken.DICT:
70 | BencodeMap bencodeMap = new BencodeMap();
71 | stack.addLast(bencodeMap);
72 | isReadMap = true;
73 | return bencodeMap;
74 | case BencodeToken.INT:
75 | return readIntegerAndCheckEnd(byteFromBuffer(), BencodeToken.END);
76 | case BencodeToken.LIST:
77 | BencodeList bencodeList = new BencodeList();
78 | stack.addLast(bencodeList);
79 | isReadMap = false;
80 | return bencodeList;
81 | default:
82 | return new BencodeByteArray(readBytesAndCheckEnd(firstByte));
83 | }
84 | }
85 |
86 | /**
87 | * 只有map的key是String
88 | * 在bencode中是 数字(较短,所以可以用int) + 冒号 + byte[] 类型;
89 | * 编码格式几乎都是ascii,不过为了保险依然使用utf-8
90 | *
91 | * @return String ,如果结束,则返回null
92 | */
93 | private String readStringOrEnd() {
94 | byte b = byteFromBuffer();
95 | if (b == BencodeToken.END) {
96 | onEnd();
97 | return null;
98 | } else {
99 | return new String(readBytesAndCheckEnd(b), Bencode.getCharset());
100 | }
101 | }
102 |
103 | private byte[] readBytesAndCheckEnd(byte firstByte) {
104 | BencodeInteger length = readIntegerAndCheckEnd(firstByte, BencodeToken.SPLIT);
105 | return byteArrayFromBuffer(length.intValue());
106 | }
107 |
108 | private BencodeInteger readIntegerAndCheckEnd(byte firstByte, byte endChar) {
109 | char c = '-';
110 | if (!Character.isDigit(firstByte) && firstByte != c) {
111 | throw ex("there must be digit or char '-' ", null);
112 | }
113 | char[] chars = new char[16];
114 | int index = 1;
115 | chars[0] = (char) firstByte;
116 | while (true) {
117 | byte codePoint = byteFromBuffer();
118 | if (!Character.isDigit(codePoint)) {
119 | if (codePoint != endChar) {
120 | throw ex("there must be digit", null);
121 | }
122 | break;
123 | }
124 | // 扩容
125 | if (chars.length <= index) {
126 | char[] temp = new char[chars.length * 2];
127 | System.arraycopy(chars, 0, temp, 0, chars.length);
128 | chars = temp;
129 | }
130 | chars[index] = (char) codePoint;
131 | index++;
132 | }
133 | return new BencodeInteger(new String(chars, 0, index));
134 | }
135 |
136 | BencodeException ex(String msg, Throwable ex) {
137 | String m = "at position " + buff.position() + ":" + msg;
138 | return new BencodeException(rootMap, m, ex);
139 | }
140 |
141 | byte byteFromBuffer() {
142 | return buff.get();
143 | }
144 |
145 |
146 | byte[] byteArrayFromBuffer(int length) {
147 | byte[] bytes = new byte[length];
148 | buff.get(bytes);
149 | return bytes;
150 | }
151 |
152 | }
153 |
--------------------------------------------------------------------------------
/server/src/main/java/cc/aguesuka/btfind/dht/handler/chain/JoinDhtChain.java:
--------------------------------------------------------------------------------
1 | package cc.aguesuka.btfind.dht.handler.chain;
2 |
3 | import cc.aguesuka.bencode.BencodeMap;
4 | import cc.aguesuka.btfind.dht.KrpcToken;
5 | import cc.aguesuka.btfind.dht.beans.KrpcMessage;
6 | import cc.aguesuka.btfind.dht.handler.IDhtHandlerChain;
7 | import cc.aguesuka.btfind.util.DhtServerConfig;
8 | import cc.aguesuka.btfind.util.record.ActionEnum;
9 | import cc.aguesuka.btfind.util.record.ActionRecord;
10 | import lombok.extern.slf4j.Slf4j;
11 | import org.springframework.beans.factory.annotation.Autowired;
12 | import org.springframework.stereotype.Component;
13 |
14 | import java.net.InetAddress;
15 | import java.net.InetSocketAddress;
16 | import java.net.SocketAddress;
17 | import java.net.UnknownHostException;
18 | import java.nio.ByteBuffer;
19 | import java.security.SecureRandom;
20 | import java.util.*;
21 |
22 | /**
23 | * @author :aguesuka
24 | * 2019/9/12 13:37
25 | */
26 | @Slf4j
27 | @Component
28 | public class JoinDhtChain implements IDhtHandlerChain {
29 |
30 | private final Queue waitSend = new LinkedList<>();
31 | private final Set blackList = new HashSet<>();
32 | private final Set allNode = new HashSet<>();
33 | private final Queue newNode = new LinkedList<>();
34 | private final Queue goodNode = new LinkedList<>();
35 |
36 | private final ActionRecord record;
37 | private final Random random = new SecureRandom();
38 | private final byte[] target = new byte[20];
39 | private final DhtServerConfig config;
40 | private long lastTime;
41 |
42 | @Autowired
43 | public JoinDhtChain(ActionRecord record, DhtServerConfig config) {
44 | this.record = record;
45 | this.config = config;
46 | }
47 |
48 | private void addQueue(Queue queue, E element) {
49 | if (queue.size() >= config.getJoinDhtMaxSize()) {
50 | queue.remove();
51 | }
52 | queue.add(element);
53 | }
54 |
55 | private void moveUntil(Collection target, Queue src, int count) {
56 | while (target.size() < count && !src.isEmpty()) {
57 | target.add(src.remove());
58 | }
59 | }
60 |
61 | private void updateQueue() {
62 | int count = config.getJoinDhtCount();
63 | moveUntil(waitSend, newNode, count);
64 | if (waitSend.size() > 0 && allNode.size() <= config.getJoinDhtMaxSize()) {
65 | return;
66 | }
67 | allNode.clear();
68 | moveUntil(waitSend, goodNode, count);
69 | if (waitSend.size() > 0) {
70 | record.doRecord(ActionEnum.JOIN_DHT_CLEAR);
71 | return;
72 | }
73 | record.doRecord(ActionEnum.JOIN_DHT_RESTART);
74 | waitSend.addAll(config.getBootstrapNodes());
75 | }
76 |
77 | private void update() {
78 | if (System.currentTimeMillis() - lastTime > config.getJoinDhtInterval()) {
79 | updateQueue();
80 | lastTime = System.currentTimeMillis();
81 | record.doRecord(ActionEnum.JOIN_DHT_INTERVAL);
82 | }
83 | }
84 |
85 | private BencodeMap findNodeQuery() {
86 | BencodeMap result = new BencodeMap();
87 | byte[] transaction = new byte[4];
88 | random.nextBytes(transaction);
89 | result.putByteArray(KrpcToken.TRANSACTION, transaction);
90 | result.putString(KrpcToken.TYPE, KrpcToken.QUERY);
91 | result.putString(KrpcToken.QUERY, KrpcToken.FIND_NODE);
92 | BencodeMap argumentsMap = new BencodeMap();
93 | argumentsMap.putByteArray(KrpcToken.ID, config.getSelfNodeId());
94 | random.nextBytes(target);
95 | argumentsMap.putByteArray(KrpcToken.TARGET, target);
96 | result.put(KrpcToken.ARGUMENTS_MAP, argumentsMap);
97 | return result;
98 | }
99 |
100 | @Override
101 | public int weights() {
102 | return 0;
103 | }
104 |
105 | @Override
106 | public void onResponse(KrpcMessage response) {
107 | BencodeMap responseMap = response.getMessage().getBencodeMap(KrpcToken.RESPONSES_MAP);
108 | byte[] senderId = responseMap.getByteArray(KrpcToken.ID);
109 | if (Arrays.equals(senderId, config.getSelfNodeId())) {
110 | blackList.add(response.getAddress());
111 | return;
112 | }
113 | addQueue(goodNode, response.getAddress());
114 |
115 | byte[] nodes = response.nodes();
116 | byte[] id = new byte[KrpcToken.ID_LENGTH];
117 | byte[] ip = new byte[4];
118 | if (nodes == null) {
119 | blackList.add(response.getAddress());
120 | return;
121 | }
122 | ByteBuffer nodesBuffer = ByteBuffer.wrap(nodes);
123 | while (nodesBuffer.hasRemaining()) {
124 | nodesBuffer.get(id);
125 | nodesBuffer.get(ip);
126 | int port = nodesBuffer.getShort() & 0xffff;
127 | try {
128 | InetSocketAddress address = new InetSocketAddress(InetAddress.getByAddress(ip), port);
129 | if (allNode.contains(address) || blackList.contains(address)) {
130 | record.doRecord(ActionEnum.DHT_GET_REPEAT_ADDRESS);
131 | return;
132 | }
133 | record.doRecord(ActionEnum.DHT_GET_NEW_ADDRESS);
134 |
135 | allNode.add(address);
136 | addQueue(newNode, address);
137 | } catch (UnknownHostException e) {
138 | throw new RuntimeException(e);
139 | }
140 | }
141 | }
142 |
143 | @Override
144 | public KrpcMessage getMessage() {
145 | return new KrpcMessage(findNodeQuery(), waitSend.remove());
146 | }
147 |
148 | @Override
149 | public boolean isWriteAble() {
150 | update();
151 | return !waitSend.isEmpty();
152 | }
153 | }
154 |
--------------------------------------------------------------------------------
/server/src/main/java/cc/aguesuka/btfind/dht/handler/DhtHandler.java:
--------------------------------------------------------------------------------
1 | package cc.aguesuka.btfind.dht.handler;
2 |
3 | import cc.aguesuka.bencode.Bencode;
4 | import cc.aguesuka.bencode.BencodeMap;
5 | import cc.aguesuka.btfind.connection.NioHandler;
6 | import cc.aguesuka.btfind.dht.KrpcToken;
7 | import cc.aguesuka.btfind.dht.beans.KrpcMessage;
8 | import cc.aguesuka.btfind.util.record.ActionEnum;
9 | import cc.aguesuka.btfind.util.record.ActionRecord;
10 | import lombok.extern.slf4j.Slf4j;
11 | import org.springframework.beans.factory.annotation.Autowired;
12 | import org.springframework.context.ApplicationContext;
13 | import org.springframework.stereotype.Component;
14 |
15 | import javax.annotation.PostConstruct;
16 | import java.io.IOException;
17 | import java.net.SocketAddress;
18 | import java.nio.ByteBuffer;
19 | import java.nio.channels.DatagramChannel;
20 | import java.nio.channels.SelectionKey;
21 | import java.util.Collection;
22 | import java.util.Comparator;
23 | import java.util.List;
24 | import java.util.stream.Collectors;
25 |
26 | /**
27 | * @author :aguesuka
28 | * 2019/9/10 14:22
29 | */
30 | @Slf4j
31 | @Component
32 | public class DhtHandler implements NioHandler {
33 | private List handlerChains;
34 | private Collection queryChains;
35 | private Collection unknownChains;
36 | private final ApplicationContext applicationContext;
37 | private final ByteBuffer buffer = ByteBuffer.allocate(1024 * 64);
38 | private final ActionRecord record;
39 |
40 | @Autowired
41 | public DhtHandler(ApplicationContext applicationContext,
42 | ActionRecord record) {
43 | this.applicationContext = applicationContext;
44 | this.record = record;
45 | }
46 |
47 | /**
48 | * init chains
49 | */
50 | @PostConstruct
51 | public void init() {
52 | handlerChains = applicationContext.getBeansOfType(IDhtHandlerChain.class).values()
53 | .stream()
54 | .filter(IBaseDhtChain::enable)
55 | .sorted(Comparator.comparingInt(IDhtHandlerChain::weights).reversed())
56 | .collect(Collectors.toList());
57 | queryChains = applicationContext.getBeansOfType(IDhtQueryChain.class).values()
58 | .stream()
59 | .filter(IBaseDhtChain::enable)
60 | .collect(Collectors.toList());
61 | unknownChains = applicationContext.getBeansOfType(IDhtUnknownChain.class).values()
62 | .stream()
63 | .filter(IBaseDhtChain::enable)
64 | .collect(Collectors.toList());
65 | }
66 |
67 | private KrpcMessage readMessage(SelectionKey key) {
68 | try {
69 | DatagramChannel channel = (DatagramChannel) key.channel();
70 | buffer.clear();
71 | SocketAddress address = channel.receive(buffer);
72 | buffer.flip();
73 | if (address == null) {
74 | return null;
75 | }
76 | BencodeMap message = Bencode.parse(buffer);
77 | if (buffer.hasRemaining()) {
78 | return null;
79 | }
80 |
81 | KrpcMessage krpcMessage = new KrpcMessage(message, address);
82 | log.info("recv message `{}`", krpcMessage);
83 | record.doRecord(ActionEnum.DHT_RECV_LENGTH, buffer.position() / 1024.0);
84 | return krpcMessage;
85 | } catch (IOException | RuntimeException e) {
86 | log.warn(e.getMessage(), e);
87 | return null;
88 | }
89 | }
90 |
91 | private void sendMessage(SelectionKey key, KrpcMessage message) {
92 | try {
93 | DatagramChannel channel = (DatagramChannel) key.channel();
94 | buffer.clear();
95 | message.getMessage().writeToBuffer(buffer);
96 | buffer.flip();
97 | // fixme DatagramChannelImpl#send0 throw exception
98 | channel.send(buffer, message.getAddress());
99 | if (buffer.hasRemaining()) {
100 | throw new DhtHandlerException("send message fail:buff has remaining");
101 | }
102 | log.info("send message `{}`", message);
103 | record.doRecord(ActionEnum.DHT_SEND_SUCCESS);
104 | record.doRecord(ActionEnum.DHT_SEND_LENGTH, buffer.position() / 1024.0);
105 | } catch (IOException e) {
106 | record.doRecord(ActionEnum.DHT_SEND_FAIL);
107 | log.warn(e.getMessage(), e);
108 | }
109 | }
110 |
111 | @Override
112 | public void doHandler(SelectionKey key) {
113 | try {
114 | if (key.isReadable()) {
115 | KrpcMessage krpcMessage = readMessage(key);
116 | if (krpcMessage != null) {
117 | onReadMessage(krpcMessage);
118 | }
119 | } else if (key.isWritable()) {
120 | for (IDhtHandlerChain chain : handlerChains) {
121 | if (chain.isWriteAble()) {
122 | KrpcMessage message = chain.getMessage();
123 | sendMessage(key, message);
124 | break;
125 | }
126 | }
127 | }
128 | } finally {
129 | // update Operation-set
130 | if (handlerChains.stream().anyMatch(IDhtHandlerChain::isWriteAble)) {
131 | key.interestOps(SelectionKey.OP_WRITE | SelectionKey.OP_READ);
132 | } else {
133 | key.interestOps(SelectionKey.OP_READ);
134 | }
135 | }
136 | }
137 |
138 | private void onReadMessage(KrpcMessage message) {
139 | String type = message.type();
140 | if (type == null) {
141 | unknownChains.forEach(c -> c.onUnknownType(message));
142 | return;
143 | }
144 | switch (type) {
145 | default:
146 | unknownChains.forEach(c -> c.onUnknownType(message));
147 | break;
148 | case KrpcToken.TYPE_ERROR:
149 | unknownChains.forEach(c -> c.onRecvError(message));
150 | break;
151 | case KrpcToken.TYPE_QUERY:
152 | onQuery(message);
153 | queryChains.forEach(c -> c.onQuery(message));
154 | break;
155 | case KrpcToken.TYPE_RESPONSE:
156 | handlerChains.forEach(c -> c.onResponse(message));
157 | break;
158 | }
159 | }
160 |
161 | private void onQuery(KrpcMessage message) {
162 | String queryType = message.queryType();
163 | switch (queryType) {
164 | default:
165 | unknownChains.forEach(c -> c.onUnknownTypeQuery(message));
166 | break;
167 | case KrpcToken.PING:
168 | queryChains.forEach(c -> c.onPing(message));
169 | break;
170 | case KrpcToken.FIND_NODE:
171 | queryChains.forEach(c -> c.onFindNodes(message));
172 | break;
173 | case KrpcToken.GET_PEERS:
174 | queryChains.forEach(c -> c.onGetPeer(message));
175 | break;
176 | case KrpcToken.ANNOUNCE_PEER:
177 | queryChains.forEach(c -> c.onAnnouncePeer(message));
178 | break;
179 | }
180 | }
181 | }
182 |
--------------------------------------------------------------------------------