diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -64,7 +64,7 @@ ### Known Issue -* `dprox` has fairly large memory footprint at the moment. Over 100MB for current `dnsmasq-china-list`. +* `dprox` has fairly large memory footprint at the moment. About 85MB for current `dnsmasq-china-list`. ### License diff --git a/dprox.cabal b/dprox.cabal --- a/dprox.cabal +++ b/dprox.cabal @@ -1,13 +1,13 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.33.0. +-- This file has been generated from package.yaml by hpack version 0.34.4. -- -- see: https://github.com/sol/hpack -- --- hash: 5ad793b14e38bbe4cb7643ee582e2ba817f2912ee33cd2f21a68e5c3d20dadcf +-- hash: 595c7426c43140ed95c8e5a736712c4e64ca373cca9ba084e938e04150a516cf name: dprox -version: 0.2.0 +version: 0.3.0 synopsis: a lightweight DNS proxy server description: Please see the README on GitHub at category: DNS @@ -46,17 +46,17 @@ attoparsec >=0.13 , base >=4.12 && <5 , bytestring >=0.10 + , bytestring-trie >=0.2.4 , containers >=0.6 , dns >=3.0.4 , hashable >=1.2 , iproute >=1.7 , network >=2.8 , optparse-applicative >=0.14 - , psqueues >=0.2 + , psqueues >=0.2.3 , streaming-commons >=0.2 , time >=1.8 , unix >=2.7 - , unordered-containers >=0.2 if flag(static) ghc-options: -optl-static default-language: Haskell2010 @@ -74,6 +74,7 @@ attoparsec >=0.13 , base >=4.12 && <5 , bytestring >=0.10 + , bytestring-trie >=0.2.4 , containers >=0.6 , dns >=3.0.4 , hashable >=1.2 @@ -81,11 +82,10 @@ , iproute >=1.7 , network >=2.8 , optparse-applicative >=0.14 - , psqueues >=0.2 + , psqueues >=0.2.3 , streaming-commons >=0.2 , time >=1.8 , unix >=2.7 - , unordered-containers >=0.2 if flag(static) ghc-options: -optl-static default-language: Haskell2010 diff --git a/src/DomainRoute.hs b/src/DomainRoute.hs --- a/src/DomainRoute.hs +++ b/src/DomainRoute.hs @@ -1,7 +1,7 @@ -- SPDX-License-Identifier: BSD-3-Clause -- -- Copyright (C) 2019 Bin Jin. All Rights Reserved. -{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} module DomainRoute ( DomainRoute , newDomainRoute @@ -9,69 +9,29 @@ , getDomainRouteByPrefix ) where -import qualified Data.ByteString.Char8 as BS8 -import qualified Data.ByteString.Short as SBS -import Data.Function (on) -import Data.HashMap.Strict (HashMap) -import qualified Data.HashMap.Strict as HashMap -import Data.List (groupBy, sortOn, span) +import qualified Data.ByteString.Char8 as BS +import Data.Char (toLower) +import qualified Data.Trie as T +import qualified Data.Trie.Convenience as T import qualified Network.DNS as DNS -type DomainComp = SBS.ShortByteString - -data DomainRoute a = DomainRoute - { currentNode :: Maybe a - , children :: Maybe (HashMap DomainComp (DomainRoute a)) - } deriving (Eq, Show) - -instance Functor DomainRoute where - fmap f DomainRoute{..} = DomainRoute (fmap f currentNode) ((fmap.fmap.fmap) f children) +newtype DomainRoute a = DomainRoute (T.Trie a) + deriving (Eq, Show, Functor, Foldable) -instance Foldable DomainRoute where - foldMap f DomainRoute{..} = foldMap f currentNode `mappend` (foldMap.foldMap.foldMap) f children +normalize :: DNS.Domain -> BS.ByteString +normalize d + | BS.null d = BS.empty + | otherwise = BS.reverse (BS.cons '.' (BS.map toLower d)) newDomainRoute :: (a -> a -> a) -> [(DNS.Domain, a)] -> DomainRoute a -newDomainRoute merge servers = go sortedServers - where - fmapFst f (x, y) = (f x, y) - sortedServers = sortOn fst $ map (fmapFst getDomainComponents) servers - - go [] = DomainRoute Nothing Nothing - go lst = DomainRoute node childs - where - (prefix, suffix) = span (null . fst) lst - grouped = groupBy ((==) `on` (head.fst)) suffix - - node = if null prefix then Nothing else Just (foldr1 merge (map snd prefix)) - subtrees = [ (comp, go subtree) - | child <- grouped - , let comp = head (fst (head child)) - , let subtree = map (fmapFst tail) child - ] - childs = if null subtrees then Nothing else Just (HashMap.fromList subtrees) - -queryDomainRoute :: Bool -> DomainRoute a -> DNS.Domain -> Maybe a -queryDomainRoute allowPrefix dr domain = go dr comps +newDomainRoute merge ds = DomainRoute (T.fromListWithL' merge ds') where - comps = getDomainComponents domain - - placePrefix _ b | not allowPrefix = b - placePrefix a Nothing = a - placePrefix _ b = b - - go DomainRoute{..} [] = currentNode - go (DomainRoute node Nothing) _ = placePrefix node Nothing - go (DomainRoute node (Just child)) (x:xs) = placePrefix node $ case HashMap.lookup x child of - Nothing -> Nothing - Just subtree -> go subtree xs + ds' = [(normalize domain, record) | (domain, record) <- ds] getDomainRouteExact :: DomainRoute a -> DNS.Domain -> Maybe a -getDomainRouteExact = queryDomainRoute False +getDomainRouteExact (DomainRoute trie) d = T.lookup (normalize d) trie getDomainRouteByPrefix :: DomainRoute a -> DNS.Domain -> Maybe a -getDomainRouteByPrefix = queryDomainRoute True - -getDomainComponents :: DNS.Domain -> [DomainComp] -getDomainComponents domain = reverse $ filter (not . SBS.null) comps - where - comps = map SBS.toShort $ BS8.split '.' (DNS.normalizeCase domain) +getDomainRouteByPrefix (DomainRoute trie) d = case T.match trie (normalize d) of + Nothing -> Nothing + Just (_, res, _) -> Just res diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -17,6 +17,7 @@ build ds = [(domain, query domain) | domain <- ds] domains = ["cn", "chn.net", "red.com", "black.com", "www.red.com", "a.b.c.d.red.com"] dr1 = newDomainRoute const (build domains) + dr0 = newDomainRoute const (build ("" : domains)) describe "newDomainRoute" $ do let dr2 = newDomainRoute const (build domains ++ build domains) @@ -48,5 +49,9 @@ it "handles negative samples" $ forM_ ["", "net", "com", "chn2.net"] $ \domain -> getDomainRouteByPrefix dr1 domain `shouldBe` Nothing + it "handles root domain fallbacks" $ forM_ ["", "net", "com", "chn2.net"] $ \domain -> + getDomainRouteByPrefix dr0 domain `shouldBe` Just 0 + it "handles the longest match" $ forM_ [("c.d.red.com", query "red.com"), ("www.a.b.c.d.red.com", query "a.b.c.d.red.com")] $ \(domain, ans) -> getDomainRouteByPrefix dr1 domain `shouldBe` Just ans +