TestDnsProtocolUdpCodecFactory Class — netty Architecture

Source Code

resolver-dns/src/test/java/io/netty/resolver/dns/TestDnsServer.java lines 156–236

    private final class TestDnsProtocolUdpCodecFactory implements ProtocolCodecFactory {
        private final DnsMessageEncoder encoder = new DnsMessageEncoder();
        private final TestAAAARecordEncoder recordEncoder = new TestAAAARecordEncoder();
        private final RecordType dropRecordType;

        TestDnsProtocolUdpCodecFactory(RecordType dropRecordType) {
            this.dropRecordType = dropRecordType;
        }

        @Override
        public ProtocolEncoder getEncoder(IoSession session) {
            return new DnsUdpEncoder() {

                @Override
                public void encode(IoSession session, Object message, ProtocolEncoderOutput out) {
                    IoBuffer buf = IoBuffer.allocate(4096);
                    DnsMessage dnsMessage = filterMessage((DnsMessage) message);
                    if (dnsMessage != null) {
                        encoder.encode(buf, dnsMessage);

                        encodeAAAA(dnsMessage.getAnswerRecords(), buf);
                        encodeAAAA(dnsMessage.getAuthorityRecords(), buf);
                        encodeAAAA(dnsMessage.getAdditionalRecords(), buf);
                        buf.flip();

                        out.write(buf);
                    }
                }
            };
        }

        private void encodeAAAA(List<ResourceRecord> records, IoBuffer out) {
            for (ResourceRecord record : records) {
                // This is a hack to allow to also test for AAAA resolution as DnsMessageEncoder
                // does not support it and it is hard to extend, because the interesting methods
                // are private...
                // In case of RecordType.AAAA we need to encode the RecordType by ourselves.
                if (record.getRecordType() == RecordType.AAAA) {
                    try {
                        recordEncoder.put(out, record);
                    } catch (IOException e) {
                        // Should never happen
                        throw new IllegalStateException(e);
                    }
                }
            }
        }

        @Override
        public ProtocolDecoder getDecoder(IoSession session) {
            return new ProtocolDecoderAdapter() {
                private final DnsMessageDecoder decoder = new DnsMessageDecoder();

                @Override
                public void decode(IoSession session, IoBuffer in, ProtocolDecoderOutput out) throws IOException {
                    DnsMessage message = decoder.decode(in);
                    if (dropRecordType != null) {
                        for (QuestionRecord record: message.getQuestionRecords()) {
                            if (record.getRecordType() == dropRecordType) {
                                return;
                            }
                        }
                    }
                    out.write(message);
                }
            };
        }

        private final class TestAAAARecordEncoder extends ResourceRecordEncoder {

            @Override
            protected void putResourceRecordData(IoBuffer ioBuffer, ResourceRecord resourceRecord) {
                byte[] bytes = BYTES.get(resourceRecord.get(DnsAttribute.IP_ADDRESS));
                if (bytes == null) {
                    throw new IllegalStateException(resourceRecord.get(DnsAttribute.IP_ADDRESS));
                }
                // encode the ::1
                ioBuffer.put(bytes);
            }
        }
    }