{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

module Network.QUIC.Types.CID (
    CID (..),
    myCIDLength,
    newCID,
    fromCID,
    toCID,
    makeCID,
    unpackCID,
    nonZeroLengthCID,
    StatelessResetToken (..),
    fromStatelessResetToken,
    makeGenStatelessReset,
    PathData (..),
    newPathData,
    CIDInfo,
    newCIDInfo,
    cidInfoSeq,
    cidInfoCID,
    cidInfoSRT,
) where

import Codec.Serialise
import Crypto.Hash
import Crypto.KDF.HKDF
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Short as Short
import GHC.Generics
import Network.Socket (SockAddr)
import System.Random (getStdRandom, uniformByteString)

import Network.QUIC.Imports

myCIDLength :: Int
myCIDLength :: Int
myCIDLength = Int
8

-- | A type for conneciton ID.
newtype CID = CID Bytes deriving (CID -> CID -> Bool
(CID -> CID -> Bool) -> (CID -> CID -> Bool) -> Eq CID
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CID -> CID -> Bool
== :: CID -> CID -> Bool
$c/= :: CID -> CID -> Bool
/= :: CID -> CID -> Bool
Eq, Eq CID
Eq CID =>
(CID -> CID -> Ordering)
-> (CID -> CID -> Bool)
-> (CID -> CID -> Bool)
-> (CID -> CID -> Bool)
-> (CID -> CID -> Bool)
-> (CID -> CID -> CID)
-> (CID -> CID -> CID)
-> Ord CID
CID -> CID -> Bool
CID -> CID -> Ordering
CID -> CID -> CID
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: CID -> CID -> Ordering
compare :: CID -> CID -> Ordering
$c< :: CID -> CID -> Bool
< :: CID -> CID -> Bool
$c<= :: CID -> CID -> Bool
<= :: CID -> CID -> Bool
$c> :: CID -> CID -> Bool
> :: CID -> CID -> Bool
$c>= :: CID -> CID -> Bool
>= :: CID -> CID -> Bool
$cmax :: CID -> CID -> CID
max :: CID -> CID -> CID
$cmin :: CID -> CID -> CID
min :: CID -> CID -> CID
Ord, (forall x. CID -> Rep CID x)
-> (forall x. Rep CID x -> CID) -> Generic CID
forall x. Rep CID x -> CID
forall x. CID -> Rep CID x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CID -> Rep CID x
from :: forall x. CID -> Rep CID x
$cto :: forall x. Rep CID x -> CID
to :: forall x. Rep CID x -> CID
Generic)

instance Serialise CID

instance Show CID where
    show :: CID -> String
show (CID Bytes
cid) = Bytes -> String
shortToString (Bytes -> Bytes
enc16s Bytes
cid)

newCID :: IO CID
newCID :: IO CID
newCID = Bytes -> CID
CID (Bytes -> CID) -> IO Bytes -> IO CID
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO Bytes
getRandomBytes Int
myCIDLength

toCID :: ByteString -> CID
toCID :: ByteString -> CID
toCID = Bytes -> CID
CID (Bytes -> CID) -> (ByteString -> Bytes) -> ByteString -> CID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bytes
Short.toShort

-- | Converting a connection ID.
fromCID :: CID -> ByteString
fromCID :: CID -> ByteString
fromCID (CID Bytes
sbs) = Bytes -> ByteString
Short.fromShort Bytes
sbs

makeCID :: ShortByteString -> CID
makeCID :: Bytes -> CID
makeCID = Bytes -> CID
CID

unpackCID :: CID -> (ShortByteString, Word8)
unpackCID :: CID -> (Bytes, Word8)
unpackCID (CID Bytes
sbs) = (Bytes
sbs, Word8
len)
  where
    len :: Word8
len = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Bytes -> Int
Short.length Bytes
sbs

nonZeroLengthCID :: CID -> SockAddr -> CID
nonZeroLengthCID :: CID -> SockAddr -> CID
nonZeroLengthCID (CID Bytes
"") SockAddr
sa = ByteString -> CID
toCID (ByteString -> CID) -> ByteString -> CID
forall a b. (a -> b) -> a -> b
$ String -> ByteString
C8.pack (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ SockAddr -> String
forall a. Show a => a -> String
show SockAddr
sa
nonZeroLengthCID CID
x SockAddr
_ = CID
x

-- 16 bytes
newtype StatelessResetToken = StatelessResetToken Bytes deriving (StatelessResetToken -> StatelessResetToken -> Bool
(StatelessResetToken -> StatelessResetToken -> Bool)
-> (StatelessResetToken -> StatelessResetToken -> Bool)
-> Eq StatelessResetToken
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StatelessResetToken -> StatelessResetToken -> Bool
== :: StatelessResetToken -> StatelessResetToken -> Bool
$c/= :: StatelessResetToken -> StatelessResetToken -> Bool
/= :: StatelessResetToken -> StatelessResetToken -> Bool
Eq, Eq StatelessResetToken
Eq StatelessResetToken =>
(StatelessResetToken -> StatelessResetToken -> Ordering)
-> (StatelessResetToken -> StatelessResetToken -> Bool)
-> (StatelessResetToken -> StatelessResetToken -> Bool)
-> (StatelessResetToken -> StatelessResetToken -> Bool)
-> (StatelessResetToken -> StatelessResetToken -> Bool)
-> (StatelessResetToken
    -> StatelessResetToken -> StatelessResetToken)
-> (StatelessResetToken
    -> StatelessResetToken -> StatelessResetToken)
-> Ord StatelessResetToken
StatelessResetToken -> StatelessResetToken -> Bool
StatelessResetToken -> StatelessResetToken -> Ordering
StatelessResetToken -> StatelessResetToken -> StatelessResetToken
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: StatelessResetToken -> StatelessResetToken -> Ordering
compare :: StatelessResetToken -> StatelessResetToken -> Ordering
$c< :: StatelessResetToken -> StatelessResetToken -> Bool
< :: StatelessResetToken -> StatelessResetToken -> Bool
$c<= :: StatelessResetToken -> StatelessResetToken -> Bool
<= :: StatelessResetToken -> StatelessResetToken -> Bool
$c> :: StatelessResetToken -> StatelessResetToken -> Bool
> :: StatelessResetToken -> StatelessResetToken -> Bool
$c>= :: StatelessResetToken -> StatelessResetToken -> Bool
>= :: StatelessResetToken -> StatelessResetToken -> Bool
$cmax :: StatelessResetToken -> StatelessResetToken -> StatelessResetToken
max :: StatelessResetToken -> StatelessResetToken -> StatelessResetToken
$cmin :: StatelessResetToken -> StatelessResetToken -> StatelessResetToken
min :: StatelessResetToken -> StatelessResetToken -> StatelessResetToken
Ord)

instance Show StatelessResetToken where
    show :: StatelessResetToken -> String
show (StatelessResetToken Bytes
srt) = Bytes -> String
shortToString (Bytes -> Bytes
enc16s Bytes
srt)

fromStatelessResetToken :: StatelessResetToken -> ByteString
fromStatelessResetToken :: StatelessResetToken -> ByteString
fromStatelessResetToken (StatelessResetToken Bytes
srt) = Bytes -> ByteString
Short.fromShort Bytes
srt

makeGenStatelessReset :: IO (CID -> StatelessResetToken)
makeGenStatelessReset :: IO (CID -> StatelessResetToken)
makeGenStatelessReset = do
    ByteString
salt <- (StdGen -> (ByteString, StdGen)) -> IO ByteString
forall (m :: * -> *) a. MonadIO m => (StdGen -> (a, StdGen)) -> m a
getStdRandom ((StdGen -> (ByteString, StdGen)) -> IO ByteString)
-> (StdGen -> (ByteString, StdGen)) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> StdGen -> (ByteString, StdGen)
forall g. RandomGen g => Int -> g -> (ByteString, g)
uniformByteString Int
20
    ByteString
ikm <- (StdGen -> (ByteString, StdGen)) -> IO ByteString
forall (m :: * -> *) a. MonadIO m => (StdGen -> (a, StdGen)) -> m a
getStdRandom ((StdGen -> (ByteString, StdGen)) -> IO ByteString)
-> (StdGen -> (ByteString, StdGen)) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> StdGen -> (ByteString, StdGen)
forall g. RandomGen g => Int -> g -> (ByteString, g)
uniformByteString Int
20
    let prk :: PRK SHA256
prk = ByteString -> ByteString -> PRK SHA256
forall a salt ikm.
(HashAlgorithm a, ByteArrayAccess salt, ByteArrayAccess ikm) =>
salt -> ikm -> PRK a
extract ByteString
salt ByteString
ikm :: PRK SHA256
        makeStatelessReset :: CID -> StatelessResetToken
makeStatelessReset CID
dcid = Bytes -> StatelessResetToken
StatelessResetToken (Bytes -> StatelessResetToken) -> Bytes -> StatelessResetToken
forall a b. (a -> b) -> a -> b
$ ByteString -> Bytes
Short.toShort (ByteString -> Bytes) -> ByteString -> Bytes
forall a b. (a -> b) -> a -> b
$ PRK SHA256 -> ByteString -> Int -> ByteString
forall a info out.
(HashAlgorithm a, ByteArrayAccess info, ByteArray out) =>
PRK a -> info -> Int -> out
expand PRK SHA256
prk (CID -> ByteString
fromCID CID
dcid) Int
16
    (CID -> StatelessResetToken) -> IO (CID -> StatelessResetToken)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return CID -> StatelessResetToken
makeStatelessReset

-- 8 bytes
newtype PathData = PathData Bytes deriving (PathData -> PathData -> Bool
(PathData -> PathData -> Bool)
-> (PathData -> PathData -> Bool) -> Eq PathData
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PathData -> PathData -> Bool
== :: PathData -> PathData -> Bool
$c/= :: PathData -> PathData -> Bool
/= :: PathData -> PathData -> Bool
Eq, Int -> PathData -> ShowS
[PathData] -> ShowS
PathData -> String
(Int -> PathData -> ShowS)
-> (PathData -> String) -> ([PathData] -> ShowS) -> Show PathData
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PathData -> ShowS
showsPrec :: Int -> PathData -> ShowS
$cshow :: PathData -> String
show :: PathData -> String
$cshowList :: [PathData] -> ShowS
showList :: [PathData] -> ShowS
Show)

newPathData :: IO PathData
newPathData :: IO PathData
newPathData = Bytes -> PathData
PathData (Bytes -> PathData) -> IO Bytes -> IO PathData
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO Bytes
getRandomBytes Int
8

data CIDInfo = CIDInfo
    { CIDInfo -> Int
cidInfoSeq :: Int
    , CIDInfo -> CID
cidInfoCID :: CID
    , CIDInfo -> StatelessResetToken
cidInfoSRT :: StatelessResetToken
    }
    deriving (CIDInfo -> CIDInfo -> Bool
(CIDInfo -> CIDInfo -> Bool)
-> (CIDInfo -> CIDInfo -> Bool) -> Eq CIDInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CIDInfo -> CIDInfo -> Bool
== :: CIDInfo -> CIDInfo -> Bool
$c/= :: CIDInfo -> CIDInfo -> Bool
/= :: CIDInfo -> CIDInfo -> Bool
Eq, Eq CIDInfo
Eq CIDInfo =>
(CIDInfo -> CIDInfo -> Ordering)
-> (CIDInfo -> CIDInfo -> Bool)
-> (CIDInfo -> CIDInfo -> Bool)
-> (CIDInfo -> CIDInfo -> Bool)
-> (CIDInfo -> CIDInfo -> Bool)
-> (CIDInfo -> CIDInfo -> CIDInfo)
-> (CIDInfo -> CIDInfo -> CIDInfo)
-> Ord CIDInfo
CIDInfo -> CIDInfo -> Bool
CIDInfo -> CIDInfo -> Ordering
CIDInfo -> CIDInfo -> CIDInfo
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: CIDInfo -> CIDInfo -> Ordering
compare :: CIDInfo -> CIDInfo -> Ordering
$c< :: CIDInfo -> CIDInfo -> Bool
< :: CIDInfo -> CIDInfo -> Bool
$c<= :: CIDInfo -> CIDInfo -> Bool
<= :: CIDInfo -> CIDInfo -> Bool
$c> :: CIDInfo -> CIDInfo -> Bool
> :: CIDInfo -> CIDInfo -> Bool
$c>= :: CIDInfo -> CIDInfo -> Bool
>= :: CIDInfo -> CIDInfo -> Bool
$cmax :: CIDInfo -> CIDInfo -> CIDInfo
max :: CIDInfo -> CIDInfo -> CIDInfo
$cmin :: CIDInfo -> CIDInfo -> CIDInfo
min :: CIDInfo -> CIDInfo -> CIDInfo
Ord, Int -> CIDInfo -> ShowS
[CIDInfo] -> ShowS
CIDInfo -> String
(Int -> CIDInfo -> ShowS)
-> (CIDInfo -> String) -> ([CIDInfo] -> ShowS) -> Show CIDInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CIDInfo -> ShowS
showsPrec :: Int -> CIDInfo -> ShowS
$cshow :: CIDInfo -> String
show :: CIDInfo -> String
$cshowList :: [CIDInfo] -> ShowS
showList :: [CIDInfo] -> ShowS
Show)

newCIDInfo :: Int -> CID -> StatelessResetToken -> CIDInfo
newCIDInfo :: Int -> CID -> StatelessResetToken -> CIDInfo
newCIDInfo Int
n CID
cid StatelessResetToken
srt =
    CIDInfo
        { cidInfoSeq :: Int
cidInfoSeq = Int
n
        , cidInfoCID :: CID
cidInfoCID = CID
cid
        , cidInfoSRT :: StatelessResetToken
cidInfoSRT = StatelessResetToken
srt
        }