content: {-# LANGUAGE OverloadedStrings #-}
  
  module Network.DNS.IO (
      -- * Receiving DNS messages
      receive
    , receiveFrom
    , receiveVC
      -- * Sending pre-encoded messages
    , send
    , sendTo
    , sendVC
    , sendAll
      -- ** Encoding queries for transmission
    , encodeQuestion
    , encodeVC
      -- ** Creating query response messages
    , responseA
    , responseAAAA
    ) where
  
  import qualified Control.Exception as E
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Builder as BB
  import qualified Data.ByteString.Char8 as BS
  import qualified Data.ByteString.Lazy.Char8 as LBS
  import Data.IP (IPv4, IPv6)
  import Time.System (timeCurrent)
  import Time.Types (Elapsed(..), Seconds(..))
  import Network.Socket (Socket, SockAddr)
  import Network.Socket.ByteString (recv, recvFrom)
  import qualified Network.Socket.ByteString as Socket
  import System.IO.Error
  
  import Network.DNS.Decode (decodeAt)
  import Network.DNS.Encode (encode)
  import Network.DNS.Imports
  import Network.DNS.Types.Internal
  
  ----------------------------------------------------------------
  
  -- | Receive and decode a single 'DNSMessage' from a UDP 'Socket', throwing away
  -- the client address.  Messages longer than 'maxUdpSize' are silently
  -- truncated, but this should not occur in practice, since we cap the advertised
  -- EDNS UDP buffer size limit at the same value.  A 'DNSError' is raised if I/O
  -- or message decoding fails.
  --
  receive :: Socket -> IO DNSMessage
  receive sock = do
      let bufsiz = fromIntegral maxUdpSize
      bs <- recv sock bufsiz `E.catch` \e -> E.throwIO $ NetworkFailure e
      Elapsed (Seconds now) <- timeCurrent
      case decodeAt now bs of
          Left  e   -> E.throwIO e
          Right msg -> return msg
  
  -- | Receive and decode a single 'DNSMessage' from a UDP 'Socket'.  Messages
  -- longer than 'maxUdpSize' are silently truncated, but this should not occur
  -- in practice, since we cap the advertised EDNS UDP buffer size limit at the
  -- same value.  A 'DNSError' is raised if I/O or message decoding fails.
  --
  receiveFrom :: Socket -> IO (DNSMessage, SockAddr)
  receiveFrom sock = do
      let bufsiz = fromIntegral maxUdpSize
      (bs, client) <- recvFrom sock bufsiz `E.catch` \e -> E.throwIO $ NetworkFailure e
      Elapsed (Seconds now) <- timeCurrent
      case decodeAt now bs of
          Left  e   -> E.throwIO e
          Right msg -> return (msg, client)
  
  -- | Receive and decode a single 'DNSMesage' from a virtual-circuit (TCP).  It
  -- is up to the caller to implement any desired timeout. An 'DNSError' is
  -- raised if I/O or message decoding fails.
  --
  receiveVC :: Socket -> IO DNSMessage
  receiveVC sock = do
      len <- toLen <$> recvDNS sock 2
      bs <- recvDNS sock len
      Elapsed (Seconds now) <- timeCurrent
      case decodeAt now bs of
          Left e    -> E.throwIO e
          Right msg -> return msg
    where
      toLen bs = case B.unpack bs of
          [hi, lo] -> 256 * (fromIntegral hi) + (fromIntegral lo)
          _        -> 0              -- never reached
  
  recvDNS :: Socket -> Int -> IO ByteString
  recvDNS sock len = recv1 `E.catch` \e -> E.throwIO $ NetworkFailure e
    where
      recv1 = do
          bs1 <- recvCore len
          if BS.length bs1 == len then
              return bs1
            else do
              loop bs1
      loop bs0 = do
          let left = len - BS.length bs0
          bs1 <- recvCore left
          let bs = bs0 `BS.append` bs1
          if BS.length bs == len then
              return bs
            else
              loop bs
      eofE = mkIOError eofErrorType "connection terminated" Nothing Nothing
      recvCore len0 = do
          bs <- recv sock len0
          if bs == "" then
              E.throwIO eofE
            else
              return bs
  
  ----------------------------------------------------------------
  
  -- | Send an encoded 'DNSMessage' datagram over UDP.  The message length is
  -- implicit in the size of the UDP datagram.  With TCP you must use 'sendVC',
  -- because TCP does not have message boundaries, and each message needs to be
  -- prepended with an explicit length.  The socket must be explicitly connected
  -- to the destination nameserver.
  --
  send :: Socket -> ByteString -> IO ()
  send = (void .). Socket.send
  {-# INLINE send #-}
  
  -- | Send an encoded 'DNSMessage' datagram over UDP to a given address.  The
  -- message length is implicit in the size of the UDP datagram.  With TCP you
  -- must use 'sendVC', because TCP does not have message boundaries, and each
  -- message needs to be prepended with an explicit length.
  --
  sendTo :: Socket -> ByteString -> SockAddr -> IO ()
  sendTo sock str addr = Socket.sendTo sock str addr >> return ()
  {-# INLINE sendTo #-}
  
  -- | Send a single encoded 'DNSMessage' over TCP.  An explicit length is
  -- prepended to the encoded buffer before transmission.  If you want to
  -- send a batch of multiple encoded messages back-to-back over a single
  -- TCP connection, and then loop to collect the results, use 'encodeVC'
  -- to prefix each message with a length, and then use 'sendAll' to send
  -- a concatenated batch of the resulting encapsulated messages.
  --
  sendVC :: Socket -> ByteString -> IO ()
  sendVC = (. encodeVC). sendAll
  {-# INLINE sendVC #-}
  
  -- | Send one or more encoded 'DNSMessage' buffers over TCP, each allready
  -- encapsulated with an explicit length prefix (perhaps via 'encodeVC') and
  -- then concatenated into a single buffer.  DO NOT use 'sendAll' with UDP.
  --
  sendAll :: Socket -> BS.ByteString -> IO ()
  sendAll = Socket.sendAll
  {-# INLINE sendAll #-}
  
  -- | The encoded 'DNSMessage' has the specified request ID.  The default values
  -- of the RD, AD, CD and DO flag bits, as well as various EDNS features, can be
  -- adjusted via the 'QueryControls' parameter.
  --
  -- The caller is responsible for generating the ID via a securely seeded
  -- CSPRNG.
  --
  encodeQuestion :: Identifier     -- ^ Crypto random request id
                  -> Question      -- ^ Query name and type
                  -> QueryControls -- ^ Query flag and EDNS overrides
                  -> ByteString
  encodeQuestion idt q ctls = encode $ makeQuery idt q ctls
  
  -- | Encapsulate an encoded 'DNSMessage' buffer for transmission over a TCP
  -- virtual circuit.  With TCP the buffer needs to start with an explicit
  -- length (the length is implicit with UDP).
  --
  encodeVC :: ByteString -> ByteString
  encodeVC legacyQuery =
      let len = LBS.toStrict . BB.toLazyByteString $ BB.int16BE $ fromIntegral $ BS.length legacyQuery
      in len <> legacyQuery
  {-# INLINE encodeVC #-}
  
  ----------------------------------------------------------------
  
  -- | Compose a response with a single IPv4 RRset.  If the query
  -- had an EDNS pseudo-header, a suitable EDNS pseudo-header must
  -- be added to the response message, or else a 'FormatErr' response
  -- must be sent.  The response TTL defaults to 300 seconds, and
  -- should be updated (to the same value across all the RRs) if some
  -- other TTL value is more appropriate.
  --
  responseA :: Identifier -> Question -> [IPv4] -> DNSMessage
  responseA idt q ips = makeResponse idt q as
    where
      dom = qname q
      as  = ResourceRecord dom A classIN 300 . RD_A <$> ips
  
  -- | Compose a response with a single IPv6 RRset.  If the query
  -- had an EDNS pseudo-header, a suitable EDNS pseudo-header must
  -- be added to the response message, or else a 'FormatErr' response
  -- must be sent.  The response TTL defaults to 300 seconds, and
  -- should be updated (to the same value across all the RRs) if some
  -- other TTL value is more appropriate.
  --
  responseAAAA :: Identifier -> Question -> [IPv6] -> DNSMessage
  responseAAAA idt q ips = makeResponse idt q as
    where
      dom = qname q
      as  = ResourceRecord dom AAAA classIN 300 . RD_AAAA <$> ips
  