{-# 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 :: Socket -> IO DNSMessage
receive Socket
sock = do
    let bufsiz :: Int
bufsiz = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
maxUdpSize
    bs <- Socket -> Int -> IO ByteString
recv Socket
sock Int
bufsiz IO ByteString -> (IOException -> IO ByteString) -> IO ByteString
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \IOException
e -> DNSError -> IO ByteString
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (DNSError -> IO ByteString) -> DNSError -> IO ByteString
forall a b. (a -> b) -> a -> b
$ IOException -> DNSError
NetworkFailure IOException
e
    Elapsed (Seconds now) <- timeCurrent
    case decodeAt now bs of
        Left  DNSError
e   -> DNSError -> IO DNSMessage
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO DNSError
e
        Right DNSMessage
msg -> DNSMessage -> IO DNSMessage
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return DNSMessage
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 :: Socket -> IO (DNSMessage, SockAddr)
receiveFrom Socket
sock = do
    let bufsiz :: Int
bufsiz = Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
maxUdpSize
    (bs, client) <- Socket -> Int -> IO (ByteString, SockAddr)
recvFrom Socket
sock Int
bufsiz IO (ByteString, SockAddr)
-> (IOException -> IO (ByteString, SockAddr))
-> IO (ByteString, SockAddr)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \IOException
e -> DNSError -> IO (ByteString, SockAddr)
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (DNSError -> IO (ByteString, SockAddr))
-> DNSError -> IO (ByteString, SockAddr)
forall a b. (a -> b) -> a -> b
$ IOException -> DNSError
NetworkFailure IOException
e
    Elapsed (Seconds now) <- timeCurrent
    case decodeAt now bs of
        Left  DNSError
e   -> DNSError -> IO (DNSMessage, SockAddr)
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO DNSError
e
        Right DNSMessage
msg -> (DNSMessage, SockAddr) -> IO (DNSMessage, SockAddr)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (DNSMessage
msg, SockAddr
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 :: Socket -> IO DNSMessage
receiveVC Socket
sock = do
    len <- ByteString -> Int
forall {a}. Num a => ByteString -> a
toLen (ByteString -> Int) -> IO ByteString -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Socket -> Int -> IO ByteString
recvDNS Socket
sock Int
2
    bs <- recvDNS sock len
    Elapsed (Seconds now) <- timeCurrent
    case decodeAt now bs of
        Left DNSError
e    -> DNSError -> IO DNSMessage
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO DNSError
e
        Right DNSMessage
msg -> DNSMessage -> IO DNSMessage
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return DNSMessage
msg
  where
    toLen :: ByteString -> a
toLen ByteString
bs = case ByteString -> [Word8]
B.unpack ByteString
bs of
        [Word8
hi, Word8
lo] -> a
256 a -> a -> a
forall a. Num a => a -> a -> a
* (Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
hi) a -> a -> a
forall a. Num a => a -> a -> a
+ (Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
lo)
        [Word8]
_        -> a
0              -- never reached

recvDNS :: Socket -> Int -> IO ByteString
recvDNS :: Socket -> Int -> IO ByteString
recvDNS Socket
sock Int
len = IO ByteString
recv1 IO ByteString -> (IOException -> IO ByteString) -> IO ByteString
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \IOException
e -> DNSError -> IO ByteString
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (DNSError -> IO ByteString) -> DNSError -> IO ByteString
forall a b. (a -> b) -> a -> b
$ IOException -> DNSError
NetworkFailure IOException
e
  where
    recv1 :: IO ByteString
recv1 = do
        bs1 <- Int -> IO ByteString
recvCore Int
len
        if BS.length bs1 == len then
            return bs1
          else do
            loop bs1
    loop :: ByteString -> IO ByteString
loop ByteString
bs0 = do
        let left :: Int
left = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
bs0
        bs1 <- Int -> IO ByteString
recvCore Int
left
        let bs = ByteString
bs0 ByteString -> ByteString -> ByteString
`BS.append` ByteString
bs1
        if BS.length bs == len then
            return bs
          else
            loop bs
    eofE :: IOException
eofE = IOErrorType
-> String -> Maybe Handle -> Maybe String -> IOException
mkIOError IOErrorType
eofErrorType String
"connection terminated" Maybe Handle
forall a. Maybe a
Nothing Maybe String
forall a. Maybe a
Nothing
    recvCore :: Int -> IO ByteString
recvCore Int
len0 = do
        bs <- Socket -> Int -> IO ByteString
recv Socket
sock Int
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 :: Socket -> ByteString -> IO ()
send = (IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int -> IO ()) -> (ByteString -> IO Int) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
.)((ByteString -> IO Int) -> ByteString -> IO ())
-> (Socket -> ByteString -> IO Int)
-> Socket
-> ByteString
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Socket -> ByteString -> IO Int
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 :: Socket -> ByteString -> SockAddr -> IO ()
sendTo Socket
sock ByteString
str SockAddr
addr = Socket -> ByteString -> SockAddr -> IO Int
Socket.sendTo Socket
sock ByteString
str SockAddr
addr IO Int -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
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 :: Socket -> ByteString -> IO ()
sendVC = ((ByteString -> IO ())
-> (ByteString -> ByteString) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
encodeVC)((ByteString -> IO ()) -> ByteString -> IO ())
-> (Socket -> ByteString -> IO ()) -> Socket -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Socket -> ByteString -> IO ()
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 -> ByteString -> IO ()
sendAll = Socket -> ByteString -> IO ()
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 :: Word16 -> Question -> QueryControls -> ByteString
encodeQuestion Word16
idt Question
q QueryControls
ctls = DNSMessage -> ByteString
encode (DNSMessage -> ByteString) -> DNSMessage -> ByteString
forall a b. (a -> b) -> a -> b
$ Word16 -> Question -> QueryControls -> DNSMessage
makeQuery Word16
idt Question
q QueryControls
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 :: ByteString -> ByteString
encodeVC ByteString
legacyQuery =
    let len :: ByteString
len = LazyByteString -> ByteString
LBS.toStrict (LazyByteString -> ByteString)
-> (Builder -> LazyByteString) -> Builder -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ Int16 -> Builder
BB.int16BE (Int16 -> Builder) -> Int16 -> Builder
forall a b. (a -> b) -> a -> b
$ Int -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int16) -> Int -> Int16
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
legacyQuery
    in ByteString
len ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
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 :: Word16 -> Question -> [IPv4] -> DNSMessage
responseA Word16
idt Question
q [IPv4]
ips = Word16 -> Question -> Answers -> DNSMessage
makeResponse Word16
idt Question
q Answers
as
  where
    dom :: ByteString
dom = Question -> ByteString
qname Question
q
    as :: Answers
as  = ByteString -> TYPE -> Word16 -> TTL -> RData -> ResourceRecord
ResourceRecord ByteString
dom TYPE
A Word16
classIN TTL
300 (RData -> ResourceRecord)
-> (IPv4 -> RData) -> IPv4 -> ResourceRecord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IPv4 -> RData
RD_A (IPv4 -> ResourceRecord) -> [IPv4] -> Answers
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [IPv4]
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 :: Word16 -> Question -> [IPv6] -> DNSMessage
responseAAAA Word16
idt Question
q [IPv6]
ips = Word16 -> Question -> Answers -> DNSMessage
makeResponse Word16
idt Question
q Answers
as
  where
    dom :: ByteString
dom = Question -> ByteString
qname Question
q
    as :: Answers
as  = ByteString -> TYPE -> Word16 -> TTL -> RData -> ResourceRecord
ResourceRecord ByteString
dom TYPE
AAAA Word16
classIN TTL
300 (RData -> ResourceRecord)
-> (IPv6 -> RData) -> IPv6 -> ResourceRecord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IPv6 -> RData
RD_AAAA (IPv6 -> ResourceRecord) -> [IPv6] -> Answers
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [IPv6]
ips