gistfile1.scala

package com.airbnb.tins.api

import io.netty.handler.codec.{LengthFieldPrepender, LengthFieldBasedFrameDecoder}
import io.netty.channel._
import org.xbill.DNS.{Record, Section, Flags, Message}
import java.net.{InetSocketAddress, SocketAddress}
import io.netty.handler.logging.{LoggingHandler, LogLevel}
import io.netty.handler.codec.bytes.ByteArrayEncoder
import io.netty.bootstrap.ServerBootstrap
import io.netty.channel.nio.NioEventLoopGroup
import io.netty.channel.socket.nio.NioServerSocketChannel
import io.netty.channel.socket.SocketChannel
import io.netty.buffer.ByteBuf


private object DNS {
  val PayloadLengthFieldSize: Int = 2
  val MaxDNSFrameLength = (1 << (8 * PayloadLengthFieldSize)) + PayloadLengthFieldSize
  val LengthFieldPrepender = new LengthFieldPrepender(PayloadLengthFieldSize, false)

  object Decoder extends ChannelInboundMessageHandlerAdapter[ByteBuf] {
    def messageReceived(ctx: ChannelHandlerContext, msg: ByteBuf) {
      ctx.nextInboundMessageBuffer.add(new Message(msg.array))
    }
  }

  object Encoder extends ChannelOutboundMessageHandlerAdapter[Message] {
    def flush(ctx: ChannelHandlerContext, msg: Message) {
      ctx.channel.write(msg.toWire)
    }
  }

}

private class Serve(handler: DNSHandler,
                    listeningAddress: SocketAddress,
                    logLevel: LogLevel) {

  object ChannelHandler extends ChannelInboundMessageHandlerAdapter[Message] {
    def messageReceived(ctx: ChannelHandlerContext, req: Message) {
      new Message {
        val reqHeader = req.getHeader
        val repHeader = getHeader

        repHeader.setID(req.getHeader.getID)
        repHeader.setFlag(Flags.QR)

        val question: Record = req.getQuestion
        addRecord(question, Section.QUESTION)

        handler(question) match {
          case Success(auth, recs) => {
            if (auth) repHeader.setFlag(Flags.AA)
            for (rec <- recs) addRecord(rec, Section.ANSWER)
          }
          case Failure(reason) => repHeader.setRcode(reason)
        }
      }
    }
  }

  val byteArrayEncoder: ByteArrayEncoder = new ByteArrayEncoder()
  new ServerBootstrap().group(new NioEventLoopGroup, new NioEventLoopGroup).
    channel(classOf[NioServerSocketChannel]).
    option[java.lang.Boolean](ChannelOption.TCP_NODELAY, true).
    option[java.lang.Boolean](ChannelOption.SO_REUSEADDR, true).
    option[java.lang.Boolean](ChannelOption.SO_KEEPALIVE, true).
    option[java.lang.Integer](ChannelOption.SO_LINGER, 0).
    handler(new LoggingHandler(logLevel)).
    childHandler(new ChannelInitializer[SocketChannel] {
    def initChannel(ch: SocketChannel) {
      ch.pipeline
        .addLast("frame decoder",
        new LengthFieldBasedFrameDecoder(
          DNS.MaxDNSFrameLength,
          0, DNS.PayloadLengthFieldSize,
          0, DNS.PayloadLengthFieldSize))
        .addLast("dns decoder", DNS.Decoder)
        .addLast("frame encoder", DNS.LengthFieldPrepender)
        .addLast("byte[] encoder", byteArrayEncoder)
        .addLast("dns encoder", DNS.Encoder)
        .addLast("handler", ChannelHandler)
    }
  }).bind(listeningAddress)
}

object Serve {
  def apply(listeningAddress: SocketAddress = new InetSocketAddress(5353),
            logLevel: LogLevel = LogLevel.INFO,
            handler: DNSHandler) {
    new Serve(handler, listeningAddress, logLevel)
  }
}

gistfile2.scala

package com.airbnb.tins

import org.xbill.DNS.{Name, Record, Rcode}

package object api {
  type DNSHandler = Function1[Record, Response]

  implicit def NameToSeq(name: Name): Seq[String] =
    for (i <- name.labels - 1 to 0 by -1) yield new String(name getLabel i)
}

package api {
sealed abstract class Response
sealed abstract case class Failure(reason: Int) extends Response
object DoesNotExist extends Failure(reason = Rcode.NXDOMAIN)
object ServerFailed extends Failure(reason = Rcode.SERVFAIL)
object NotAuthorized extends Failure(reason = Rcode.NOTAUTH)
object NotImplemented extends Failure(reason = Rcode.NOTIMP)
case class Success(authoritative: Boolean = true,
                   records: Seq[Record]) extends Response
}