View Javadoc
1   /*
2    * Copyright 2015 The Netty Project
3    *
4    * The Netty Project licenses this file to you under the Apache License,
5    * version 2.0 (the "License"); you may not use this file except in compliance
6    * with the License. You may obtain a copy of the License at:
7    *
8    *   http://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13   * License for the specific language governing permissions and limitations
14   * under the License.
15   */
16  package io.netty.handler.codec.dns;
17  
18  import io.netty.buffer.ByteBuf;
19  import io.netty.channel.AddressedEnvelope;
20  import io.netty.channel.ChannelHandler;
21  import io.netty.channel.ChannelHandlerContext;
22  import io.netty.channel.socket.DatagramPacket;
23  import io.netty.handler.codec.MessageToMessageEncoder;
24  import io.netty.util.internal.UnstableApi;
25  
26  import java.net.InetSocketAddress;
27  import java.util.List;
28  
29  import static io.netty.util.internal.ObjectUtil.checkNotNull;
30  
31  /**
32   * Encodes a {@link DatagramDnsQuery} (or an {@link AddressedEnvelope} of {@link DnsQuery}} into a
33   * {@link DatagramPacket}.
34   */
35  @UnstableApi
36  @ChannelHandler.Sharable
37  public class DatagramDnsQueryEncoder extends MessageToMessageEncoder<AddressedEnvelope<DnsQuery, InetSocketAddress>> {
38  
39      private final DnsRecordEncoder recordEncoder;
40  
41      /**
42       * Creates a new encoder with {@linkplain DnsRecordEncoder#DEFAULT the default record encoder}.
43       */
44      public DatagramDnsQueryEncoder() {
45          this(DnsRecordEncoder.DEFAULT);
46      }
47  
48      /**
49       * Creates a new encoder with the specified {@code recordEncoder}.
50       */
51      public DatagramDnsQueryEncoder(DnsRecordEncoder recordEncoder) {
52          this.recordEncoder = checkNotNull(recordEncoder, "recordEncoder");
53      }
54  
55      @Override
56      protected void encode(
57          ChannelHandlerContext ctx,
58          AddressedEnvelope<DnsQuery, InetSocketAddress> in, List<Object> out) throws Exception {
59  
60          final InetSocketAddress recipient = in.recipient();
61          final DnsQuery query = in.content();
62          final ByteBuf buf = allocateBuffer(ctx, in);
63  
64          boolean success = false;
65          try {
66              encodeHeader(query, buf);
67              encodeQuestions(query, buf);
68              encodeRecords(query, DnsSection.ADDITIONAL, buf);
69              success = true;
70          } finally {
71              if (!success) {
72                  buf.release();
73              }
74          }
75  
76          out.add(new DatagramPacket(buf, recipient, null));
77      }
78  
79      /**
80       * Allocate a {@link ByteBuf} which will be used for constructing a datagram packet.
81       * Sub-classes may override this method to return a {@link ByteBuf} with a perfect matching initial capacity.
82       */
83      protected ByteBuf allocateBuffer(
84          ChannelHandlerContext ctx,
85          @SuppressWarnings("unused") AddressedEnvelope<DnsQuery, InetSocketAddress> msg) throws Exception {
86          return ctx.alloc().ioBuffer(1024);
87      }
88  
89      /**
90       * Encodes the header that is always 12 bytes long.
91       *
92       * @param query the query header being encoded
93       * @param buf   the buffer the encoded data should be written to
94       */
95      private static void encodeHeader(DnsQuery query, ByteBuf buf) {
96          buf.writeShort(query.id());
97          int flags = 0;
98          flags |= (query.opCode().byteValue() & 0xFF) << 14;
99          if (query.isRecursionDesired()) {
100             flags |= 1 << 8;
101         }
102         buf.writeShort(flags);
103         buf.writeShort(query.count(DnsSection.QUESTION));
104         buf.writeShort(0); // answerCount
105         buf.writeShort(0); // authorityResourceCount
106         buf.writeShort(query.count(DnsSection.ADDITIONAL));
107     }
108 
109     private void encodeQuestions(DnsQuery query, ByteBuf buf) throws Exception {
110         final int count = query.count(DnsSection.QUESTION);
111         for (int i = 0; i < count; i++) {
112             recordEncoder.encodeQuestion((DnsQuestion) query.recordAt(DnsSection.QUESTION, i), buf);
113         }
114     }
115 
116     private void encodeRecords(DnsQuery query, DnsSection section, ByteBuf buf) throws Exception {
117         final int count = query.count(section);
118         for (int i = 0; i < count; i++) {
119             recordEncoder.encodeRecord(query.recordAt(section, i), buf);
120         }
121     }
122 }