{-# LANGUAGE CPP #-}

module Data.ByteString.Base64.Extra
  ( encodeBase64StdUnpadded
  , decodeBase64StdUnpadded
  ) where

#if MIN_VERSION_base64(1,0,0)
import qualified Data.Base64.Types as B64
#endif
import Data.ByteString ( ByteString )
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base64 as B64
import qualified Data.ByteString.Base64.URL as B64URL
import Data.Text ( Text )
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import Prelude

encodeBase64StdUnpadded :: ByteString -> ByteString
encodeBase64StdUnpadded :: ByteString -> ByteString
encodeBase64StdUnpadded =
  -- Drop padding bytes ('=') from the end of the base64 string.
  ((Word8 -> Bool) -> ByteString -> ByteString
BS.dropWhileEnd (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x3D))
#if MIN_VERSION_base64(1,0,0)
    (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Base64 'StdPadded ByteString -> ByteString
forall (k :: Alphabet) a. Base64 k a -> a
B64.extractBase64
#endif
    (Base64 'StdPadded ByteString -> ByteString)
-> (ByteString -> Base64 'StdPadded ByteString)
-> ByteString
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Base64 'StdPadded ByteString
B64.encodeBase64'

decodeBase64StdUnpadded :: ByteString -> Either Text ByteString
decodeBase64StdUnpadded :: ByteString -> Either Text ByteString
decodeBase64StdUnpadded ByteString
b64 =
  -- The @base64@ library does not support decoding unpadded standard base64.
  -- So we're going to convert to base64url and then decode.
  let b64Url :: ByteString
      b64Url :: ByteString
b64Url = Text -> ByteString
TE.encodeUtf8 (HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"+" Text
"-" (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"/" Text
"_" (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
TE.decodeUtf8 ByteString
b64)
  in
#if MIN_VERSION_base64(1,0,0)
    ByteString -> Either Text ByteString
B64URL.decodeBase64UnpaddedUntyped ByteString
b64Url
#else
    B64URL.decodeBase64Unpadded b64Url
#endif