content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module Network.QUIC.Connection.Crypto (
      setEncryptionLevel,
      waitEncryptionLevel,
      putOffCrypto,
      --
      getCipher,
      setCipher,
      getTLSMode,
      getApplicationProtocol,
      setNegotiated,
      --
      dropSecrets,
      --
      initializeCoder,
      initializeCoder1RTT,
      updateCoder1RTT,
      getCoder,
      getProtector,
      --
      getCurrentKeyPhase,
      setCurrentKeyPhase,
  ) where
  
  import Control.Concurrent.STM
  import Network.TLS.QUIC
  
  import Network.QUIC.Connection.Misc
  import Network.QUIC.Connection.Types
  import Network.QUIC.Connector
  import Network.QUIC.Crypto
  import Network.QUIC.Imports
  import Network.QUIC.Types
  
  useFusion :: Bool
  #ifdef USE_FUSION
  useFusion = True
  #else
  useFusion = False
  #endif
  
  ----------------------------------------------------------------
  
  setEncryptionLevel :: Connection -> EncryptionLevel -> IO ()
  setEncryptionLevel Connection{..} lvl = do
      let q = connRecvQ
      atomically $ do
          writeTVar (encryptionLevel connState) lvl
          case lvl of
              HandshakeLevel -> do
                  readTVar (pendingQ ! RTT0Level) >>= mapM_ (prependRecvQ q)
                  readTVar (pendingQ ! HandshakeLevel) >>= mapM_ (prependRecvQ q)
              RTT1Level ->
                  readTVar (pendingQ ! RTT1Level) >>= mapM_ (prependRecvQ q)
              _ -> return ()
  
  putOffCrypto :: Connection -> EncryptionLevel -> ReceivedPacket -> IO ()
  putOffCrypto Connection{..} lvl rpkt =
      atomically $ modifyTVar' (pendingQ ! lvl) (rpkt :)
  
  waitEncryptionLevel :: Connection -> EncryptionLevel -> IO ()
  waitEncryptionLevel Connection{..} lvl = atomically $ do
      l <- readTVar $ encryptionLevel connState
      check (l >= lvl)
  
  ----------------------------------------------------------------
  
  getCipher :: Connection -> EncryptionLevel -> IO Cipher
  getCipher Connection{..} lvl = readArray ciphers lvl
  
  setCipher :: Connection -> EncryptionLevel -> Cipher -> IO ()
  setCipher Connection{..} lvl cipher = writeArray ciphers lvl cipher
  
  ----------------------------------------------------------------
  
  getTLSMode :: Connection -> IO HandshakeMode13
  getTLSMode Connection{..} = tlsHandshakeMode <$> readIORef negotiated
  
  getApplicationProtocol :: Connection -> IO (Maybe NegotiatedProtocol)
  getApplicationProtocol Connection{..} = applicationProtocol <$> readIORef negotiated
  
  setNegotiated
      :: Connection
      -> HandshakeMode13
      -> Maybe NegotiatedProtocol
      -> ApplicationSecretInfo
      -> IO ()
  setNegotiated Connection{..} mode mproto appSecInf =
      writeIORef
          negotiated
          Negotiated
              { tlsHandshakeMode = mode
              , applicationProtocol = mproto
              , applicationSecretInfo = appSecInf
              }
  
  ----------------------------------------------------------------
  
  dropSecrets :: Connection -> EncryptionLevel -> IO ()
  dropSecrets Connection{..} lvl = do
      writeArray coders lvl initialCoder
      writeArray protectors lvl initialProtector
  
  ----------------------------------------------------------------
  
  initializeCoder :: Connection -> EncryptionLevel -> TrafficSecrets a -> IO ()
  initializeCoder conn lvl sec = do
      ver <-
          if lvl == RTT0Level
              then return $ getOriginalVersion conn
              else getVersion conn
      cipher <- getCipher conn lvl
      avail <- isFusionAvailable
      (coder, protector) <-
          if useFusion && avail
              then genFusionCoder (isClient conn) ver cipher sec
              else genNiteCoder (isClient conn) ver cipher sec
      writeArray (coders conn) lvl coder
      writeArray (protectors conn) lvl protector
  
  initializeCoder1RTT :: Connection -> TrafficSecrets ApplicationSecret -> IO ()
  initializeCoder1RTT conn sec = do
      ver <- getVersion conn
      cipher <- getCipher conn RTT1Level
      avail <- isFusionAvailable
      (coder, protector) <-
          if useFusion && avail
              then genFusionCoder (isClient conn) ver cipher sec
              else genNiteCoder (isClient conn) ver cipher sec
      let coder1 = Coder1RTT coder sec
      writeArray (coders1RTT conn) False coder1
      writeArray (protectors conn) RTT1Level protector
      updateCoder1RTT conn True
  
  updateCoder1RTT :: Connection -> Bool -> IO ()
  updateCoder1RTT conn nextPhase = do
      ver <- getVersion conn
      cipher <- getCipher conn RTT1Level
      Coder1RTT coder secN <- readArray (coders1RTT conn) (not nextPhase)
      let secN1 = updateSecret ver cipher secN
      avail <- isFusionAvailable
      coderN1 <-
          if useFusion && avail
              then genFusionCoder1RTT (isClient conn) ver cipher secN1 coder
              else genNiteCoder1RTT (isClient conn) ver cipher secN1 coder
      let nextCoder = Coder1RTT coderN1 secN1
      writeArray (coders1RTT conn) nextPhase nextCoder
  
  updateSecret
      :: Version
      -> Cipher
      -> TrafficSecrets ApplicationSecret
      -> TrafficSecrets ApplicationSecret
  updateSecret ver cipher (ClientTrafficSecret cN, ServerTrafficSecret sN) = secN1
    where
      Secret cN1 = nextSecret ver cipher $ Secret cN
      Secret sN1 = nextSecret ver cipher $ Secret sN
      secN1 = (ClientTrafficSecret cN1, ServerTrafficSecret sN1)
  
  genFusionCoder
      :: Bool -> Version -> Cipher -> TrafficSecrets a -> IO (Coder, Protector)
  genFusionCoder cli ver cipher (ClientTrafficSecret c, ServerTrafficSecret s) = do
      fctxt <- fusionNewContext
      fctxr <- fusionNewContext
      fusionSetup cipher fctxt txPayloadKey txPayloadIV
      fusionSetup cipher fctxr rxPayloadKey rxPayloadIV
      supp <- fusionSetupSupplement cipher txHeaderKey
      let coder =
              Coder
                  { encrypt = fusionEncrypt fctxt supp
                  , decrypt = fusionDecrypt fctxr
                  , supplement = Just supp
                  }
      let protector =
              Protector
                  { setSample = fusionSetSample supp
                  , getMask = fusionGetMask supp
                  , unprotect = unp
                  }
      return (coder, protector)
    where
      txSecret
          | cli = Secret c
          | otherwise = Secret s
      rxSecret
          | cli = Secret s
          | otherwise = Secret c
      txPayloadKey = aeadKey ver cipher txSecret
      txPayloadIV = initialVector ver cipher txSecret
      txHeaderKey = headerProtectionKey ver cipher txSecret
      rxPayloadKey = aeadKey ver cipher rxSecret
      rxPayloadIV = initialVector ver cipher rxSecret
      rxHeaderKey = headerProtectionKey ver cipher rxSecret
      unp = protectionMask cipher rxHeaderKey
  
  genNiteCoder
      :: Bool -> Version -> Cipher -> TrafficSecrets a -> IO (Coder, Protector)
  genNiteCoder cli ver cipher (ClientTrafficSecret c, ServerTrafficSecret s) = do
      let enc = makeNiteEncrypt cipher txPayloadKey txPayloadIV
          dec = makeNiteDecrypt cipher rxPayloadKey rxPayloadIV
      (set, get) <- makeNiteProtector cipher txHeaderKey
      let coder =
              Coder
                  { encrypt = enc
                  , decrypt = dec
                  , supplement = Nothing
                  }
      let protector =
              Protector
                  { setSample = set
                  , getMask = get
                  , unprotect = unp
                  }
      return (coder, protector)
    where
      txSecret
          | cli = Secret c
          | otherwise = Secret s
      rxSecret
          | cli = Secret s
          | otherwise = Secret c
      txPayloadKey = aeadKey ver cipher txSecret
      txPayloadIV = initialVector ver cipher txSecret
      txHeaderKey = headerProtectionKey ver cipher txSecret
      rxPayloadKey = aeadKey ver cipher rxSecret
      rxPayloadIV = initialVector ver cipher rxSecret
      rxHeaderKey = headerProtectionKey ver cipher rxSecret
      unp = protectionMask cipher rxHeaderKey
  
  genFusionCoder1RTT
      :: Bool -> Version -> Cipher -> TrafficSecrets a -> Coder -> IO Coder
  genFusionCoder1RTT cli ver cipher (ClientTrafficSecret c, ServerTrafficSecret s) oldcoder = do
      fctxt <- fusionNewContext
      fctxr <- fusionNewContext
      fusionSetup cipher fctxt txPayloadKey txPayloadIV
      fusionSetup cipher fctxr rxPayloadKey rxPayloadIV
      let supp = fromJust $ supplement oldcoder
      let coder =
              Coder
                  { encrypt = fusionEncrypt fctxt supp
                  , decrypt = fusionDecrypt fctxr
                  , supplement = Just supp
                  }
      return coder
    where
      txSecret
          | cli = Secret c
          | otherwise = Secret s
      rxSecret
          | cli = Secret s
          | otherwise = Secret c
      txPayloadKey = aeadKey ver cipher txSecret
      txPayloadIV = initialVector ver cipher txSecret
      rxPayloadKey = aeadKey ver cipher rxSecret
      rxPayloadIV = initialVector ver cipher rxSecret
  
  genNiteCoder1RTT
      :: Bool -> Version -> Cipher -> TrafficSecrets a -> Coder -> IO Coder
  genNiteCoder1RTT cli ver cipher (ClientTrafficSecret c, ServerTrafficSecret s) _oldcoder = do
      let enc = makeNiteEncrypt cipher txPayloadKey txPayloadIV
          dec = makeNiteDecrypt cipher rxPayloadKey rxPayloadIV
      let coder =
              Coder
                  { encrypt = enc
                  , decrypt = dec
                  , supplement = Nothing
                  }
      return coder
    where
      txSecret
          | cli = Secret c
          | otherwise = Secret s
      rxSecret
          | cli = Secret s
          | otherwise = Secret c
      txPayloadKey = aeadKey ver cipher txSecret
      txPayloadIV = initialVector ver cipher txSecret
      rxPayloadKey = aeadKey ver cipher rxSecret
      rxPayloadIV = initialVector ver cipher rxSecret
  
  getCoder :: Connection -> EncryptionLevel -> Bool -> IO Coder
  getCoder conn RTT1Level k = coder1RTT <$> readArray (coders1RTT conn) k
  getCoder conn lvl _ = readArray (coders conn) lvl
  
  getProtector :: Connection -> EncryptionLevel -> IO Protector
  getProtector conn lvl = readArray (protectors conn) lvl
  
  ----------------------------------------------------------------
  
  getCurrentKeyPhase :: Connection -> IO (Bool, PacketNumber)
  getCurrentKeyPhase Connection{..} = readIORef currentKeyPhase
  
  setCurrentKeyPhase :: Connection -> Bool -> PacketNumber -> IO ()
  setCurrentKeyPhase Connection{..} k pn = writeIORef currentKeyPhase (k, pn)
  