dprox-0.4.0: src/Main.hs
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Copyright (C) 2019 Bin Jin. All Rights Reserved.
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
module Main where
import Control.Concurrent (forkIO, threadDelay)
import Control.Exception (SomeException, handle)
import Control.Monad (forM, forever, join, void)
import Data.ByteString (ByteString)
import Data.Foldable qualified as F
import Data.Hashable (Hashable (..))
import Data.IP qualified as IP
import Data.IP.RouteTable qualified as IP
import Data.Map qualified as M
import Data.Maybe (fromMaybe, isJust, maybeToList)
import Data.Set qualified as S
import Data.Streaming.Network (bindPortUDP)
import Data.Version (showVersion)
import Network.DNS qualified as DNS
import Network.Socket.ByteString (recvFrom, sendTo)
import System.Posix.User
(UserEntry (..), getUserEntryForName, setUserID)
import Config
import DomainRoute
import Log
import LRU
import Paths_dprox
instance Hashable DNS.TYPE where
hashWithSalt s = hashWithSalt s . DNS.fromTYPE
type Resolver = DNS.Domain -> DNS.TYPE -> IO (Either DNS.DNSError [DNS.RData])
type CachedResolver = DNS.Domain -> DNS.TYPE -> IO (Either DNS.DNSError (DNS.TTL, [DNS.RData]))
processQuery :: CachedResolver -> DNS.Question -> IO [DNS.ResourceRecord]
processQuery resolver (DNS.Question qd qt) = handle handler $ do
res <- resolver qd qt
case res of
Left _ -> return []
Right (ttl, r) -> return (map (wrapper ttl) r)
where
handler :: SomeException -> IO [DNS.ResourceRecord]
handler _ = return []
wrapper :: DNS.TTL -> DNS.RData -> DNS.ResourceRecord
wrapper ttl rdata = DNS.ResourceRecord qd (getType rdata) DNS.classIN ttl rdata
getType :: DNS.RData -> DNS.TYPE
getType DNS.RD_A{} = DNS.A
getType DNS.RD_NS{} = DNS.NS
getType DNS.RD_CNAME{} = DNS.CNAME
getType DNS.RD_SOA{} = DNS.SOA
getType DNS.RD_NULL{} = DNS.NULL
getType DNS.RD_PTR{} = DNS.PTR
getType DNS.RD_MX{} = DNS.MX
getType DNS.RD_TXT{} = DNS.TXT
getType DNS.RD_AAAA{} = DNS.AAAA
getType _ = qt
parseDNS :: ByteString -> Either DNS.DNSError DNS.DNSMessage
parseDNS bs = do
q <- DNS.decode bs
if DNS.qOrR (DNS.flags (DNS.header q)) == DNS.QR_Query
then return q
else Left DNS.FormatError
processDNS :: CachedResolver -> DNS.DNSMessage -> IO (Either DNS.DNSError ByteString)
processDNS resolver DNS.DNSMessage{ DNS.header = hd, DNS.question = q } = do
rrs <- mapM (processQuery resolver) q
let hd0 = DNS.header DNS.defaultResponse
resp = DNS.defaultResponse {
DNS.header = hd0 { DNS.identifier = DNS.identifier hd },
DNS.question = q,
DNS.answer = concat rrs
}
return (Right (DNS.encode resp))
handleServer :: DomainRoute (Maybe Resolver) -> Resolver
handleServer route qd qt = case resolver of
Nothing -> return (Left DNS.NameError)
Just resolver' -> resolver' qd qt
where
resolver = join (getDomainRouteByPrefix route qd)
handleAddressAndHosts :: DomainRoute [IP] -> DomainRoute [IP] -> Resolver -> Resolver
handleAddressAndHosts address hosts resolver qd qt =
if null ips then resolver qd qt else return (Right userDefined)
where
ips1 = fromMaybe [] $ getDomainRouteByPrefix address qd
ips2 = fromMaybe [] $ getDomainRouteExact hosts qd
ips | null ips2 = ips1
| otherwise = ips2
ipv4 = [DNS.RD_A ipv4addr | IPv4 ipv4addr <- ips]
ipv6 = [DNS.RD_AAAA ipv6addr | IPv6 ipv6addr <- ips]
userDefined | qt == DNS.A = ipv4
| qt == DNS.AAAA = ipv6
| otherwise = []
handleBogusNX :: S.Set IP -> Resolver -> Resolver
handleBogusNX blacklist resolver qd qt =
fmap (filter (not . isBlacklisted)) <$> resolver qd qt
where
isBlacklisted (DNS.RD_A ipv4) = IPv4 ipv4 `S.member` blacklist
isBlacklisted (DNS.RD_AAAA ipv6) = IPv6 ipv6 `S.member` blacklist
isBlacklisted _ = False
handleIPSet :: [IPMask] -> IPSetMatch -> Maybe Resolver -> LRUCache DNS.Domain () -> CachedResolver -> CachedResolver
handleIPSet [] _ _ _ = id
handleIPSet _ _ Nothing _ = id
handleIPSet ipset match (Just ipsetResolver) cache = handleWithResolver
where
iproute = IP.fromList $ zip ipset $ repeat ()
inIPSet ip = isJust $ IP.lookup (IP.makeAddrRange ip 32) iproute
check NoneMatch = not . any inIPSet
check AllMatch = all inIPSet
check AnyMatch = any inIPSet
check AnyNotMatch = any (not.inIPSet)
handleWithResolver resolver qd qt@DNS.A = do
cachedInIPSet <- isJust <$> lookupCache qd cache
if cachedInIPSet
then fmap (1,) <$> ipsetResolver qd qt
else do
res <- resolver qd qt
case res of
Left _ -> return res
Right (_, rs) -> do
let ipv4s = [ipv4 | DNS.RD_A ipv4 <- rs]
if not (null ipv4s) && check match ipv4s
then do
updateCache qd () cache
fmap (1,) <$> ipsetResolver qd qt
else return res
handleWithResolver resolver qd qt = resolver qd qt
makeResolverCache :: Int -> DNS.TTL -> IO (Resolver -> CachedResolver)
makeResolverCache sz ttl | sz <= 0 = return $ \r qd qt -> fmap (ttl,) <$> r qd qt
makeResolverCache sz ttl = do
cache <- newCache sz ttl
_ <- forkIO $ forever $ do
threadDelay (fromInteger (fromIntegral ttl * 1000000 `div` 3))
purgeCache cache
let process resolver qd qt = do
let k = (qd, qt)
res <- lookupCache k cache
case res of
Just (ttl', v) -> return (Right (max ttl' 3, v))
Nothing -> do
resolved <- resolver qd qt
case resolved of
Left e -> return (Left e)
Right v -> do
updateCache k v cache
return (Right (ttl, v))
return process
setuid :: String -> IO ()
setuid user = getUserEntryForName user >>= setUserID . userID
main :: IO ()
main = getConfig >>= \(GlobalConfig{..}, conf) -> withLogger (LogStdout 4096) loglevel $ \logger -> do
logger INFO $ "dprox " <> toLogStr (showVersion version) <> " started"
let defaultPort = 53
fallbackServer = Server Nothing "8.8.8.8" Nothing
server = [ (fromMaybe "" mDomain, (ip, fromMaybe defaultPort mPort))
| Server mDomain ip mPort <- fallbackServer : conf
]
address = [(domain, [ip]) | Address domain ip <- conf]
hosts = [(domain, [ip]) | Hosts domain ip <- conf]
bogusnx = [ip | BogusNX ip <- conf]
ipset = [ipmask | IPSet ipmask <- conf]
serverRoute = newDomainRoute (const id) server
serverAddressSet = S.fromList $ F.toList serverRoute
addressRoute = newDomainRoute (++) address
hostsRoute = newDomainRoute (++) hosts
bogusnxSet = S.fromList bogusnx
ipsetServerPort = fmap (fromMaybe defaultPort) <$> ipsetServer
resolvConfs = [ (addr, rc)
| addr@(host, port) <- S.toList serverAddressSet ++ maybeToList ipsetServerPort
, host /= invalidIPAddress
, let rsinfo = if port == defaultPort
then DNS.RCHostName (show host)
else DNS.RCHostPort (show host) port
, let rc = DNS.defaultResolvConf {
DNS.resolvCache = Nothing,
DNS.resolvInfo = rsinfo
}
]
logger INFO $ "read " <> toLogStr (length server) <> " server configs"
logger INFO $ "read " <> toLogStr (length hosts) <> " hosts configs"
logger INFO $ "read " <> toLogStr (length bogusnx) <> " bogus-nxdomain configs"
logger INFO $ "read " <> toLogStr (length ipset) <> " ipset configs"
let (bport, bhost) = (fromIntegral $ fromMaybe defaultPort localPort, fromMaybe "*6" listenAddress)
sock <- bindPortUDP bport bhost
logger INFO $ "bind to " <> toLogStr (show bhost) <> ":" <> toLogStr bport
resolvSeeds <- forM resolvConfs $ \(k, v) -> do
logger INFO $ "creating resolver: " <> toLogStr (show $ fst k) <> ":" <> toLogStr (show $ snd k)
rs <- DNS.makeResolvSeed v
return (k, rs)
F.mapM_ setuid setUser
resolverCache <- makeResolverCache cacheSize cacheTTL
ipsetCache <- newCache 4096 maxBound
let processWithResolver resolver = forever $ do
(bs, addr) <- recvFrom sock (fromIntegral DNS.maxUdpSize)
forkIO $ do
case parseDNS bs of
Left err -> logger DEBUG $ "dns parse error: " <> toLogStr (show err) <> " from: " <> toLogStr (show addr)
Right q -> do
logger TRACE $ "dns request from " <> toLogStr (show addr) <> ": " <> toLogStr (show $ DNS.question q)
resp <- processDNS resolver q
case resp of
Right resp' -> void $ sendTo sock resp' addr
Left err -> do
logger DEBUG $ "dns process error: " <> toLogStr (show err) <> " from: " <> toLogStr (show addr)
createResolvers ((k,v):xs) m = DNS.withResolver v $ \rs ->
createResolvers xs (M.insert k (DNS.lookup rs) m)
createResolvers [] m = let serverRoute' = fmap (`M.lookup`m) serverRoute
ipsetResolver = join $ fmap (`M.lookup`m) ipsetServerPort
resolver = handleIPSet ipset ipsetMatch ipsetResolver ipsetCache $
resolverCache $
handleAddressAndHosts addressRoute hostsRoute $
handleBogusNX bogusnxSet $
handleServer serverRoute'
in processWithResolver resolver
createResolvers resolvSeeds M.empty