-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

-- Parts of this code, namely functions interpretResponse and orderSrvResult,
-- which were taken from http://hackage.haskell.org/package/pontarius-xmpp
-- are also licensed under the three-clause BSD license:
--
-- Copyright © 2005-2011 Dmitry Astapov
-- Copyright © 2005-2011 Pierre Kovalev
-- Copyright © 2010-2011 Mahdi Abdinejadi
-- Copyright © 2010-2013 Jon Kristensen
-- Copyright © 2011      IETF Trust
-- Copyright © 2012-2013 Philipp Balzarek
--
-- All rights reserved.
--
-- Pontarius XMPP is licensed under the three-clause BSD license.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--
-- - Redistributions of source code must retain the above copyright notice, this
--   list of conditions and the following disclaimer.
--
-- - Redistributions in binary form must reproduce the above copyright notice,
--   this list of conditions and the following disclaimer in the documentation
--   and/or other materials provided with the distribution.
--
-- - Neither the name of the Pontarius project nor the names of its contributors
--   may be used to endorse or promote products derived from this software without
--   specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR THE PONTARIUS PROJECT BE
-- LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-- CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
-- GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
-- HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
-- LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
-- OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

module Wire.Network.DNS.SRV where

import Control.Category ((>>>))
import Data.List.NonEmpty (NonEmpty (..))
import Imports
import Network.DNS (DNSError, Domain)
import System.Random (randomRIO)

data SrvEntry = SrvEntry
  { srvPriority :: !Word16,
    srvWeight :: !Word16,
    srvTarget :: !SrvTarget
  }
  deriving (Eq, Show, Ord)

data SrvTarget = SrvTarget
  { -- | the hostname on which the service is offered
    srvTargetDomain :: !Domain,
    -- | the port on which the service is offered
    srvTargetPort :: !Word16
  }
  deriving (Eq, Show, Ord)

data SrvResponse
  = SrvNotAvailable
  | SrvAvailable (NonEmpty SrvEntry)
  | SrvResponseError DNSError
  deriving (Eq, Show)

interpretResponse :: Either DNSError [(Word16, Word16, Word16, Domain)] -> SrvResponse
interpretResponse = \case
  Left err -> SrvResponseError err
  Right [] -> SrvNotAvailable
  Right [(_, _, _, ".")] -> SrvNotAvailable -- According to RFC2782
  Right (r : rs) -> SrvAvailable $ fmap toSrvEntry (r :| rs)

toSrvEntry :: (Word16, Word16, Word16, Domain) -> SrvEntry
toSrvEntry (prio, weight, port, domain) = SrvEntry prio weight (SrvTarget domain port)

-- FUTUREWORK: maybe improve sorting algorithm here? (with respect to performance and code style)
--
-- This function orders the SRV result in accordance with RFC
-- 2782. It sorts the SRV results in order of priority, and then
-- uses a random process to order the records with the same
-- priority based on their weight.
--
-- Taken from http://hackage.haskell.org/package/pontarius-xmpp (BSD3 licence) and refactored.
orderSrvResult :: [SrvEntry] -> IO [SrvEntry]
orderSrvResult =
  -- Order the result set by priority.
  sortOn srvPriority
    -- Group elements in sublists based on their priority.
    -- The result type is `[[(Word16, Word16, Word16, Domain)]]' (nested list).
    >>> groupBy ((==) `on` srvPriority)
    -- For each sublist, put records with a weight of zero first.
    >>> map (uncurry (++) . partition ((== 0) . srvWeight))
    -- Order each sublist.
    >>> mapM orderSublist
    -- Concatenate the results.
    >>> fmap concat
  where
    orderSublist :: [SrvEntry] -> IO [SrvEntry]
    orderSublist [] = pure []
    orderSublist sublist = do
      -- Compute the running sum, as well as the total sum of the sublist.
      -- Add the running sum to the SRV tuples.
      let (total, sublistWithRunning) =
            mapAccumL (\acc srv -> let acc' = acc + srvWeight srv in (acc', (srv, acc'))) 0 sublist
      -- Choose a random number between 0 and the total sum (inclusive).
      randomNumber <- randomRIO (0, total)
      -- Select the first record with its running sum greater
      -- than or equal to the random number.
      let (beginning, (firstSrv, _), end) =
            case break (\(_, running) -> randomNumber <= running) sublistWithRunning of
              (b, c : e) -> (b, c, e)
              _ -> error "orderSrvResult: no record with running sum greater than random number"
      -- Remove the running total number from the remaining elements.
      let remainingSrvs = map fst (beginning ++ end)
      -- Repeat the ordering procedure on the remaining elements.
      rest <- orderSublist remainingSrvs
      pure $ firstSrv : rest
