RatchetTree.hs

-- 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/>.

module Wire.API.MLS.RatchetTree where

import Data.Functor
import Imports
import Wire.API.MLS.Extension
import Wire.API.MLS.HPKEPublicKey
import Wire.API.MLS.LeafNode
import Wire.API.MLS.Serialisation

data RatchetTree = RatchetTree {nodes :: [Maybe RatchetTreeNode]}
  deriving (Eq, Show)

instance IsExtension RatchetTree where
  extensionType = 2

instance ParseMLS RatchetTree where
  parseMLS = RatchetTree <$> parseMLSVector @VarInt (parseMLSOptional parseMLS)

data RatchetTreeNodeTag = RatchetTreeLeafNodeTag | RatchetTreeParentNodeTag
  deriving (Eq, Ord, Show, Enum, Bounded)

instance ParseMLS RatchetTreeNodeTag where
  parseMLS = parseMLSEnum @Word8 "NodeType"

data RatchetTreeNode
  = RatchetTreeParentNode RatchetTreeParent
  | RatchetTreeLeafNode LeafNode
  deriving (Eq, Show)

instance ParseMLS RatchetTreeNode where
  parseMLS =
    parseMLS >>= \case
      RatchetTreeParentNodeTag -> RatchetTreeParentNode <$> parseMLS
      RatchetTreeLeafNodeTag -> RatchetTreeLeafNode <$> parseMLS

data RatchetTreeParent = RatchetTreeParent
  { encryptionKey :: HPKEPublicKey,
    parentHash :: ByteString,
    unmergedLeaves :: [Word32]
  }
  deriving (Eq, Show)

instance ParseMLS RatchetTreeParent where
  parseMLS =
    RatchetTreeParent
      <$> parseMLS
      <*> parseMLSBytes @VarInt
      <*> parseMLSVector @VarInt parseMLS

ratchetTreeLeaves :: RatchetTree -> [(LeafIndex, LeafNode)]
ratchetTreeLeaves =
  foldMap (traverse toNode) . zip [0 ..] . evens . (.nodes)
  where
    evens :: [a] -> [a]
    evens [] = []
    evens [x] = [x]
    evens (x : _ : xs) = (x : evens xs)

    toNode :: Maybe RatchetTreeNode -> [LeafNode]
    toNode (Just (RatchetTreeLeafNode n)) = [n]
    toNode _ = []