Base64.hs

{-# LANGUAGE OverloadedStrings #-}

module Network.DomainAuth.Pubkey.Base64 (
    decode,
) where

import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (Builder)
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy as BL
import Data.Word
import Network.DomainAuth.Utils

-- | Decoding base64 with folding white spaces.
--   'Nothing' is returned for malformed input.
--
-- >>> decode "aGVs bG8="
-- Just "hello"
-- >>> decode "aGVsbG8"
-- Nothing
-- >>> decode "aGV*bG8="
-- Nothing
-- >>> decode "=AAA"
-- Nothing
decode :: ByteString -> Maybe ByteString
decode bs = BL.toStrict . BB.toLazyByteString <$> dec (BS.filter (not . isSpace) bs)

dec :: ByteString -> Maybe Builder
dec bs
    | BS.null bs = Just empty
    | len < 4 = Nothing
    | len == 4 && c3 == cEqual && c4 == cEqual = dec1' <$> x1 <*> x2
    | len == 4 && c4 == cEqual = dec2' <$> x1 <*> x2 <*> x3
    | otherwise = (+++) <$> (dec' <$> x1 <*> x2 <*> x3 <*> x4) <*> dec bs'
  where
    len = BS.length bs
    c3 = bs !!! 2
    c4 = bs !!! 3
    x1 = fromChar (bs !!! 0)
    x2 = fromChar (bs !!! 1)
    x3 = fromChar c3
    x4 = fromChar c4
    bs' = BS.drop 4 bs

dec' :: Word8 -> Word8 -> Word8 -> Word8 -> Builder
dec' x1 x2 x3 x4 = BB.word8 d1 <> BB.word8 d2 <> BB.word8 d3
  where
    d1 = (x1 `shiftL` 2) .|. (x2 `shiftR` 4)
    d2 = ((x2 `shiftL` 4) .&. 0xF0) .|. (x3 `shiftR` 2)
    d3 = ((x3 `shiftL` 6) .&. 0xC0) .|. x4

dec1' :: Word8 -> Word8 -> Builder
dec1' x1 x2 = BB.word8 d1
  where
    d1 = (x1 `shiftL` 2) .|. (x2 `shiftR` 4)

dec2' :: Word8 -> Word8 -> Word8 -> Builder
dec2' x1 x2 x3 = BB.word8 d1 <> BB.word8 d2
  where
    d1 = (x1 `shiftL` 2) .|. (x2 `shiftR` 4)
    d2 = ((x2 `shiftL` 4) .&. 0xF0) .|. (x3 `shiftR` 2)

fromChar :: Word8 -> Maybe Word8
fromChar c
    | isUpper c = Just (c - cA)
    | isLower c = Just (c - cSmallA + 26)
    | isDigit c = Just (c - cZero + 52)
    | c == cPlus = Just 62
    | c == cSlash = Just 63
    | otherwise = Nothing