Resolver.hs

{-# LANGUAGE OverloadedStrings #-}

module Network.DomainAuth.SPF.Resolver (
    resolveSPF,
    SpfLookup (..),
    dnsLookup,
    SpfCtx,
    newSpfCtx,
) where

import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS
import Data.IORef
import Data.IP
import Data.Maybe
import Network.DNS
import Network.DomainAuth.SPF.Parser
import Network.DomainAuth.SPF.Types

----------------------------------------------------------------

-- | DNS lookup functions used by SPF.
data SpfLookup = SpfLookup
    { spfLookupTXT :: Domain -> IO (Either DNSError [ByteString])
    , spfLookupA :: Domain -> IO (Either DNSError [IPv4])
    , spfLookupAAAA :: Domain -> IO (Either DNSError [IPv6])
    , spfLookupMX :: Domain -> IO (Either DNSError [(Domain, Int)])
    }

dnsLookup :: Resolver -> SpfLookup
dnsLookup r =
    SpfLookup
        { spfLookupTXT = lookupTXT r
        , spfLookupA = lookupA r
        , spfLookupAAAA = lookupAAAA r
        , spfLookupMX = lookupMX r
        }

-- | Shared state for one SPF evaluation.
data SpfCtx = SpfCtx
    { ctxLookup :: SpfLookup
    , ctxLimit :: Int
    , ctxCount :: IORef Int
    , ctxVoid :: IORef Int
    }

newSpfCtx :: SpfLookup -> Int -> IO SpfCtx
newSpfCtx lk lim = SpfCtx lk lim <$> newIORef 0 <*> newIORef 0

-- RFC 7208 Section 4.6.4
voidLookupLimit, mxLimit :: Int
voidLookupLimit = 2
mxLimit = 10

-- | Charging one DNS-querying term against the limit.
charge :: SpfCtx -> IO ()
charge ctx = do
    n <- atomicModifyIORef' (ctxCount ctx) $ \x -> (x + 1, x + 1)
    when (n > ctxLimit ctx) $ fail "PermError"

-- | NXDOMAIN is treated as empty. Empty answers are void lookups.
checkVoid :: SpfCtx -> Either DNSError [a] -> IO [a]
checkVoid ctx r = case r of
    Left NameError -> void' >> return []
    Left _ -> fail "TempError"
    Right [] -> void' >> return []
    Right xs -> return xs
  where
    void' = do
        n <- atomicModifyIORef' (ctxVoid ctx) $ \x -> (x + 1, x + 1)
        when (n > voidLookupLimit) $ fail "PermError"

----------------------------------------------------------------

-- | 'True' for include/redirect, where no SPF record is PermError.
resolveSPF :: SpfCtx -> Bool -> Domain -> IP -> IO [IO SpfSeq]
resolveSPF ctx nested dom ip = do
    jrc <- spfLookupTXT (ctxLookup ctx) dom
    rc <- case jrc of
        Left NameError -> return []
        Left _ -> fail "TempError"
        Right rc -> return rc
    let rr = getSPFRR rc
    checkExistence rr (if nested then "PermError" else "None")
    let jrs = parseSPF rr
    checkSyntax jrs "PermError"
    let is = filterSPFWithIP ip (fromJust jrs)
    return $ map (toSpfSeq ctx dom ip) is
  where
    getSPFRR rc = case filter ("v=spf1" `BS.isPrefixOf`) rc of
        [] -> ""
        t : _ -> t
    checkSyntax rs estr = when (isNothing rs) (fail estr)
    checkExistence rr estr = when (BS.null rr) (fail estr)

----------------------------------------------------------------

filterSPFWithIP :: IP -> [SPF] -> [SPF]
filterSPFWithIP (IPv4 _) spfs = filter exceptIPv4 spfs
filterSPFWithIP (IPv6 _) spfs = filter exceptIPv6 spfs

exceptIPv4 :: SPF -> Bool
exceptIPv4 (SPF_IPv6Range _ _) = False
exceptIPv4 _ = True

exceptIPv6 :: SPF -> Bool
exceptIPv6 (SPF_IPv4Range _ _) = False
exceptIPv6 _ = True

----------------------------------------------------------------

toSpfSeq :: SpfCtx -> Domain -> IP -> SPF -> IO SpfSeq
toSpfSeq _ _ _ (SPF_IPv4Range q ipr) = return $ SS_IPv4Range q ipr
toSpfSeq _ _ _ (SPF_IPv6Range q ipr) = return $ SS_IPv6Range q ipr
toSpfSeq _ _ _ (SPF_All q) = return $ SS_All q
toSpfSeq ctx _ ip (SPF_Include q dom) = do
    charge ctx
    SS_IF_Pass q <$> resolveSPF ctx True dom ip
toSpfSeq ctx _ ip (SPF_Redirect dom) = do
    charge ctx
    SS_SpfSeq <$> resolveSPF ctx True dom ip
toSpfSeq ctx dom0 ip (SPF_MX q mdom (l4, l6)) = do
    charge ctx
    mxs <- spfLookupMX lk dom >>= checkVoid ctx
    when (length mxs > mxLimit) $ fail "PermError"
    let hosts = map fst mxs
    case ip of
        IPv4 _ ->
            SS_IPv4Ranges q . ranges l4 . concat <$> mapM (lookupHost (spfLookupA lk)) hosts
        IPv6 _ ->
            SS_IPv6Ranges q . ranges l6 . concat
                <$> mapM (lookupHost (spfLookupAAAA lk)) hosts
  where
    lk = ctxLookup ctx
    dom = fromMaybe dom0 mdom
toSpfSeq ctx dom0 ip (SPF_Address q mdom (l4, l6)) = do
    charge ctx
    case ip of
        IPv4 _ -> SS_IPv4Ranges q . ranges l4 <$> (spfLookupA lk dom >>= checkVoid ctx)
        IPv6 _ -> SS_IPv6Ranges q . ranges l6 <$> (spfLookupAAAA lk dom >>= checkVoid ctx)
  where
    lk = ctxLookup ctx
    dom = fromMaybe dom0 mdom

ranges :: Addr a => Int -> [a] -> [AddrRange a]
ranges len = map (`makeAddrRange` len)

-- | Lookups for MX hosts are not counted as void lookups.
lookupHost :: (Domain -> IO (Either DNSError [a])) -> Domain -> IO [a]
lookupHost f d = do
    r <- f d
    case r of
        Left NameError -> return []
        Left _ -> fail "TempError"
        Right xs -> return xs