{-# LANGUAGE FlexibleInstances, OverloadedStrings, CPP #-}

-- |
-- Module:      Network.Riak.Connection
-- Copyright:   (c) 2011 MailRank, Inc.
-- License:     Apache
-- Maintainer:  Tim McGilchrist <timmcgil@gmail.com>, Mark Hibberd <mark@hibberd.id.au>, Nathan Hunter <nhunter@janrain.com>
-- Stability:   experimental
-- Portability: portable
--
-- Support for REST-safe name handling.
--
-- Riak's protocol buffer (PBC) API will accept unescaped bucket,
-- link, and key names.  Its REST API does not unescape names, so it
-- is possible to use the PBC API to construct names that cannot be
-- accessed via the REST API (e.g. containing an embedded slash or
-- other URL-unsafe octet).

module Network.Riak.Escape
    (
      Escape(..)
    , unescape
    ) where

import Blaze.ByteString.Builder (Builder, fromByteString, toByteString, toLazyByteString)
import Blaze.ByteString.Builder.Word (fromWord8)
#if __GLASGOW_HASKELL__ < 710
import Control.Applicative ((<$>))
#endif
import Data.Attoparsec.ByteString as A
import Data.Bits ((.|.), (.&.), shiftL, shiftR)
import Data.ByteString (ByteString)
#if __GLASGOW_HASKELL__ < 710
import Data.Monoid (mappend, mempty)
#endif
import Data.Text (Text)
import Data.Word (Word8)
import Data.Bifunctor (second, first)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Unsafe as B
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL

-- | The class of string-like types that can be URL-escaped and
-- unescaped.
class Escape e where
    -- | URL-escape a string.
    escape :: e -> ByteString
    -- | URL-unescape a string.
    unescape' :: ByteString -> Either String e

-- | URL-unescape a string that is presumed to be properly escaped.
-- If the string is invalid, an error will be thrown that cannot be
-- caught from pure code.
unescape :: Escape e => ByteString -> e
unescape :: ByteString -> e
unescape ByteString
bs = case ByteString -> Either String e
forall e. Escape e => ByteString -> Either String e
unescape' ByteString
bs of
                Left String
err -> String -> e
forall a. HasCallStack => String -> a
error (String -> e) -> String -> e
forall a b. (a -> b) -> a -> b
$ String
"Network.Riak.Escape.unescape: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
                Right e
v  -> e
v
{-# INLINE unescape #-}

instance Escape ByteString where
    escape :: ByteString -> ByteString
escape = ByteString -> ByteString
L.toStrict (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> (ByteString -> Builder) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Builder -> Word8 -> Builder) -> Builder -> ByteString -> Builder
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl Builder -> Word8 -> Builder
escapeWord8 Builder
forall a. Monoid a => a
mempty
    {-# INLINE escape #-}
    unescape' :: ByteString -> Either String ByteString
unescape' = Result ByteString -> Either String ByteString
forall r. Result r -> Either String r
A.eitherResult (Result ByteString -> Either String ByteString)
-> (ByteString -> Result ByteString)
-> ByteString
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
A.parse (Builder -> ByteString
toByteString (Builder -> ByteString)
-> Parser ByteString Builder -> Parser ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Builder
unescapeBS)
    {-# INLINE unescape' #-}

instance Escape L.ByteString where
    escape :: ByteString -> ByteString
escape = ByteString -> ByteString
L.toStrict (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> (ByteString -> Builder) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Builder -> Word8 -> Builder) -> Builder -> ByteString -> Builder
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
L.foldl Builder -> Word8 -> Builder
escapeWord8 Builder
forall a. Monoid a => a
mempty
    {-# INLINE escape #-}
    unescape' :: ByteString -> Either String ByteString
unescape' = Result ByteString -> Either String ByteString
forall r. Result r -> Either String r
A.eitherResult (Result ByteString -> Either String ByteString)
-> (ByteString -> Result ByteString)
-> ByteString
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
A.parse (Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> Parser ByteString Builder -> Parser ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Builder
unescapeBS)
    {-# INLINE unescape' #-}

instance Escape Text where
    escape :: Text -> ByteString
escape = ByteString -> ByteString
forall e. Escape e => e -> ByteString
escape (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf8
    {-# INLINE escape #-}
    unescape' :: ByteString -> Either String Text
unescape' = (Either String ByteString
-> (ByteString -> Either String Text) -> Either String Text
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (UnicodeException -> String)
-> Either UnicodeException Text -> Either String Text
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first UnicodeException -> String
forall a. Show a => a -> String
show (Either UnicodeException Text -> Either String Text)
-> (ByteString -> Either UnicodeException Text)
-> ByteString
-> Either String Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either UnicodeException Text
T.decodeUtf8') (Either String ByteString -> Either String Text)
-> (ByteString -> Either String ByteString)
-> ByteString
-> Either String Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Result ByteString -> Either String ByteString
forall r. Result r -> Either String r
A.eitherResult (Result ByteString -> Either String ByteString)
-> (ByteString -> Result ByteString)
-> ByteString
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
A.parse (Builder -> ByteString
toByteString (Builder -> ByteString)
-> Parser ByteString Builder -> Parser ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Builder
unescapeBS)
    {-# INLINE unescape' #-}

instance Escape TL.Text where
    escape :: Text -> ByteString
escape = ByteString -> ByteString
forall e. Escape e => e -> ByteString
escape (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TL.encodeUtf8
    {-# INLINE escape #-}
    unescape' :: ByteString -> Either String Text
unescape' = (Either String ByteString
-> (ByteString -> Either String Text) -> Either String Text
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (UnicodeException -> String)
-> Either UnicodeException Text -> Either String Text
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first UnicodeException -> String
forall a. Show a => a -> String
show (Either UnicodeException Text -> Either String Text)
-> (ByteString -> Either UnicodeException Text)
-> ByteString
-> Either String Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either UnicodeException Text
TL.decodeUtf8') (Either String ByteString -> Either String Text)
-> (ByteString -> Either String ByteString)
-> ByteString
-> Either String Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Result ByteString -> Either String ByteString
forall r. Result r -> Either String r
A.eitherResult (Result ByteString -> Either String ByteString)
-> (ByteString -> Result ByteString)
-> ByteString
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
A.parse (Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> Parser ByteString Builder -> Parser ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Builder
unescapeBS)
    {-# INLINE unescape' #-}

instance Escape [Char] where
    escape :: String -> ByteString
escape = ByteString -> ByteString
forall e. Escape e => e -> ByteString
escape (ByteString -> ByteString)
-> (String -> ByteString) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> (String -> Text) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack
    {-# INLINE escape #-}
    unescape' :: ByteString -> Either String String
unescape' = (Text -> String) -> Either String Text -> Either String String
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Text -> String
T.unpack (Either String Text -> Either String String)
-> (ByteString -> Either String Text)
-> ByteString
-> Either String String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either String Text
forall e. Escape e => ByteString -> Either String e
unescape'
    {-# INLINE unescape' #-}

-- | URL-escape a byte from a bytestring.
escapeWord8 :: Builder -> Word8 -> Builder
escapeWord8 :: Builder -> Word8 -> Builder
escapeWord8 Builder
acc Word8
32 = Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
fromWord8 Word8
43
escapeWord8 Builder
acc Word8
i
    | Word8 -> Bool
literal Word8
i = Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
fromWord8 Word8
i
    | Bool
otherwise = Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
hex Word8
i
  where
    literal :: Word8 -> Bool
literal Word8
w = Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
97 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
122 Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
65 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
90 Bool -> Bool -> Bool
||
                Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
48 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
57 Bool -> Bool -> Bool
|| Word8
w Word8 -> ByteString -> Bool
`B.elem` ByteString
"$-.!*'(),_"
    hex :: Word8 -> Builder
hex Word8
w = Word8 -> Builder
fromWord8 Word8
37 Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
d (Word8
w Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
4) Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
d (Word8
w Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0xf)
    d :: Word8 -> Builder
d Word8
n | Word8
n Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
10    = Word8 -> Builder
fromWord8 (Word8
n Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
48)
        | Bool
otherwise = Word8 -> Builder
fromWord8 (Word8
n Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
87)
{-# INLINE escapeWord8 #-}

-- | URL-unescape' a bytestring.
unescapeBS :: Parser Builder
unescapeBS :: Parser ByteString Builder
unescapeBS = Builder -> Parser ByteString Builder
go Builder
forall a. Monoid a => a
mempty
  where
    go :: Builder -> Parser ByteString Builder
go Builder
acc  = do
      ByteString
s <- (Word8 -> Bool) -> Parser ByteString
A.takeWhile ((Word8 -> Bool) -> Parser ByteString)
-> (Word8 -> Bool) -> Parser ByteString
forall a b. (a -> b) -> a -> b
$ \Word8
w -> Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
37 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
43
      let rest :: Parser ByteString Builder
rest = do
            Word8
w <- Parser Word8
anyWord8
            if Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
43
              then Builder -> Parser ByteString Builder
go (Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` ByteString -> Builder
fromByteString ByteString
s Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` Word8 -> Builder
fromWord8 Word8
32)
              else do
                ByteString
h <- Int -> Parser ByteString
A.take Int
2
                let hex :: a -> a
hex a
b | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
48 Bool -> Bool -> Bool
&& a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
57  = a
b a -> a -> a
forall a. Num a => a -> a -> a
- a
48
                          | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
97 Bool -> Bool -> Bool
&& a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
102 = a
b a -> a -> a
forall a. Num a => a -> a -> a
- a
87
                          | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
65 Bool -> Bool -> Bool
&& a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
70  = a
b a -> a -> a
forall a. Num a => a -> a -> a
- a
55
                          | Bool
otherwise           = a
255
                    hi :: Word8
hi = Word8 -> Word8
forall a. (Ord a, Num a) => a -> a
hex (ByteString -> Int -> Word8
B.unsafeIndex ByteString
h Int
0)
                    lo :: Word8
lo = Word8 -> Word8
forall a. (Ord a, Num a) => a -> a
hex (ByteString -> Int -> Word8
B.unsafeIndex ByteString
h Int
1)
                if Word8
hi Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
lo Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
255
                  then String -> Parser ByteString Builder
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"invalid hex escape"
                  else Builder -> Parser ByteString Builder
go (Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` ByteString -> Builder
fromByteString ByteString
s Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend`
                           Word8 -> Builder
fromWord8 (Word8
lo Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. (Word8
hi Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
4)))
      Bool
done <- Parser ByteString Bool
forall t. Chunk t => Parser t Bool
atEnd
      if Bool
done
        then Builder -> Parser ByteString Builder
forall (m :: * -> *) a. Monad m => a -> m a
return (Builder
acc Builder -> Builder -> Builder
forall a. Monoid a => a -> a -> a
`mappend` ByteString -> Builder
fromByteString ByteString
s)
        else Parser ByteString Builder
rest