Adding HMAC-SHA256 authentication.

This commit is contained in:
Andrew Gibiansky 2015-03-02 15:08:08 -08:00
parent 0f39ef9729
commit b4cc01df7d
5 changed files with 141 additions and 131 deletions

View File

@ -49,7 +49,8 @@ library
transformers >=0.3,
unix >=2.6,
uuid >=1.3,
zeromq4-haskell >=0.1
zeromq4-haskell >=0.1,
SHA >=1.6
-- Example program

View File

@ -173,7 +173,7 @@ easyKernel :: (MonadIO m)
-> m ()
easyKernel profileFile config = do
prof <- liftIO $ getProfile profileFile
zmq@(Channels shellReqChan shellRepChan ctrlReqChan ctrlRepChan iopubChan) <-
zmq@(Channels shellReqChan shellRepChan ctrlReqChan ctrlRepChan iopubChan _) <-
liftIO $ serveProfile prof
execCount <- liftIO $ newMVar 0
forever $ do

View File

@ -6,18 +6,14 @@
-- `parseMessage`, which should only be used in the low-level 0MQ interface.
module IHaskell.IPython.Message.Parser (parseMessage) where
import Data.Aeson ((.:), decode, Result(..), Object)
import Control.Applicative ((<|>), (<$>), (<*>))
import Data.Aeson.Types (parse)
import Data.ByteString
import Data.Map (Map)
import Data.Text (Text)
import qualified Data.ByteString.Lazy as Lazy
import IHaskell.IPython.Types
import Debug.Trace
import Data.Aeson ((.:), decode, Result(..), Object)
import Control.Applicative ((<|>), (<$>), (<*>))
import Data.Aeson.Types (parse)
import Data.ByteString
import Data.Map (Map)
import Data.Text (Text)
import qualified Data.ByteString.Lazy as Lazy
import IHaskell.IPython.Types
type LByteString = Lazy.ByteString
@ -25,16 +21,16 @@ type LByteString = Lazy.ByteString
-- | Parse a message from its ByteString components into a Message.
parseMessage :: [ByteString] -- ^ The list of identifiers sent with the message.
-> ByteString -- ^ The header data.
-> ByteString -- ^ The parent header, which is just "{}" if there is no header.
-> ByteString -- ^ The metadata map, also "{}" for an empty map.
-> ByteString -- ^ The message content.
-> Message -- ^ A parsed message.
parseMessage idents headerData parentHeader metadata content =
-> ByteString -- ^ The header data.
-> ByteString -- ^ The parent header, which is just "{}" if there is no header.
-> ByteString -- ^ The metadata map, also "{}" for an empty map.
-> ByteString -- ^ The message content.
-> Message -- ^ A parsed message.
parseMessage idents headerData parentHeader metadata content =
let header = parseHeader idents headerData parentHeader metadata
messageType = msgType header
messageWithoutHeader = parser messageType $ Lazy.fromStrict content in
messageWithoutHeader { header = header }
messageWithoutHeader = parser messageType $ Lazy.fromStrict content
in messageWithoutHeader { header = header }
----- Module internals -----
@ -44,50 +40,50 @@ parseHeader :: [ByteString] -- ^ The list of identifiers.
-> ByteString -- ^ The parent header, or "{}" for Nothing.
-> ByteString -- ^ The metadata, or "{}" for an empty map.
-> MessageHeader -- The resulting message header.
parseHeader idents headerData parentHeader metadata = MessageHeader {
identifiers = idents,
parentHeader = parentResult,
metadata = metadataMap,
messageId = messageUUID,
sessionId = sessionUUID,
username = username,
msgType = messageType
} where
-- Decode the header data and the parent header data into JSON objects.
-- If the parent header data is absent, just have Nothing instead.
Just result = decode $ Lazy.fromStrict headerData :: Maybe Object
parentResult = if parentHeader == "{}"
parseHeader idents headerData parentHeader metadata =
MessageHeader { identifiers = idents
, parentHeader = parentResult
, metadata = metadataMap
, messageId = messageUUID
, sessionId = sessionUUID
, username = username
, msgType = messageType
}
where
-- Decode the header data and the parent header data into JSON objects.
-- If the parent header data is absent, just have Nothing instead.
Just result = decode $ Lazy.fromStrict headerData :: Maybe Object
parentResult = if parentHeader == "{}"
then Nothing
else Just $ parseHeader idents parentHeader "{}" metadata
Success (messageType, username, messageUUID, sessionUUID) = traceShow result $ flip parse result $ \obj -> do
messType <- obj .: "msg_type"
username <- obj .: "username"
message <- obj .: "msg_id"
session <- obj .: "session"
return (messType, username, message, session)
Success (messageType, username, messageUUID, sessionUUID) = flip parse result $ \obj -> do
messType <- obj .: "msg_type"
username <- obj .: "username"
message <- obj .: "msg_id"
session <- obj .: "session"
return (messType, username, message, session)
-- Get metadata as a simple map.
Just metadataMap = decode $ Lazy.fromStrict metadata :: Maybe (Map Text Text)
-- Get metadata as a simple map.
Just metadataMap = decode $ Lazy.fromStrict metadata :: Maybe (Map Text Text)
noHeader :: MessageHeader
noHeader = error "No header created"
parser :: MessageType -- ^ The message type being parsed.
-> LByteString -> Message -- The parser that converts the body into a message.
-- This message should have an undefined
-- header.
parser KernelInfoRequestMessage = kernelInfoRequestParser
parser ExecuteRequestMessage = executeRequestParser
parser CompleteRequestMessage = completeRequestParser
parser ObjectInfoRequestMessage = objectInfoRequestParser
parser ShutdownRequestMessage = shutdownRequestParser
parser InputReplyMessage = inputReplyParser
parser CommOpenMessage = commOpenParser
parser CommDataMessage = commDataParser
parser CommCloseMessage = commCloseParser
parser HistoryRequestMessage = historyRequestParser
parser other = error $ "Unknown message type " ++ show other
-> LByteString -> Message -- ^ The parser that converts the body into a message.
-- This message should have an undefined header.
parser KernelInfoRequestMessage = kernelInfoRequestParser
parser ExecuteRequestMessage = executeRequestParser
parser CompleteRequestMessage = completeRequestParser
parser ObjectInfoRequestMessage = objectInfoRequestParser
parser ShutdownRequestMessage = shutdownRequestParser
parser InputReplyMessage = inputReplyParser
parser CommOpenMessage = commOpenParser
parser CommDataMessage = commDataParser
parser CommCloseMessage = commCloseParser
parser HistoryRequestMessage = historyRequestParser
parser other = error $ "Unknown message type " ++ show other
-- | Parse a kernel info request.
-- A kernel info request has no auxiliary information, so ignore the body.
@ -101,7 +97,7 @@ kernelInfoRequestParser _ = KernelInfoRequest { header = noHeader }
-- 3. "store_history": whether to include this in history.
-- 4. "allow_stdin": whether to allow reading from stdin for this code.
executeRequestParser :: LByteString -> Message
executeRequestParser content =
executeRequestParser content =
let parser obj = do
code <- obj .: "code"
silent <- obj .: "silent"
@ -110,21 +106,20 @@ executeRequestParser content =
return (code, silent, storeHistory, allowStdin)
Just decoded = decode content
Success (code, silent, storeHistory, allowStdin) = parse parser decoded in
ExecuteRequest {
header = noHeader,
getCode = code,
getSilent = silent,
getAllowStdin = allowStdin,
getStoreHistory = storeHistory,
getUserVariables = [],
getUserExpressions = []
}
Success (code, silent, storeHistory, allowStdin) = parse parser decoded
in ExecuteRequest { header = noHeader
, getCode = code
, getSilent = silent
, getAllowStdin = allowStdin
, getStoreHistory = storeHistory
, getUserVariables = []
, getUserExpressions = []
}
requestParser parser content = parsed
where
Success parsed = parse parser decoded
Just decoded = decode content
Success parsed = parse parser decoded
Just decoded = decode content
historyRequestParser :: LByteString -> Message
historyRequestParser = requestParser $ \obj ->

View File

@ -53,25 +53,24 @@ type Port = Int
type IP = String
-- | The transport mechanism used to communicate with the IPython frontend.
data Transport
= TCP -- ^ Default transport mechanism via TCP.
deriving (Show, Read)
data Transport = TCP -- ^ Default transport mechanism via TCP.
deriving (Show, Read)
-- | A kernel profile, specifying how the kernel communicates.
data Profile = Profile {
ip :: IP, -- ^ The IP on which to listen.
transport :: Transport, -- ^ The transport mechanism.
stdinPort :: Port, -- ^ The stdin channel port.
controlPort :: Port, -- ^ The control channel port.
hbPort :: Port, -- ^ The heartbeat channel port.
shellPort :: Port, -- ^ The shell command port.
iopubPort :: Port, -- ^ The IOPub port.
key :: Text -- ^ The HMAC encryption key.
} deriving (Show, Read)
data Profile = Profile { ip :: IP -- ^ The IP on which to listen.
, transport :: Transport -- ^ The transport mechanism.
, stdinPort :: Port -- ^ The stdin channel port.
, controlPort :: Port -- ^ The control channel port.
, hbPort :: Port -- ^ The heartbeat channel port.
, shellPort :: Port -- ^ The shell command port.
, iopubPort :: Port -- ^ The IOPub port.
, signatureKey :: ByteString -- ^ The HMAC encryption key.
}
deriving (Show, Read)
-- Convert the kernel profile to and from JSON.
instance FromJSON Profile where
parseJSON (Object v) =
parseJSON (Object v) =
Profile <$> v .: "ip"
<*> v .: "transport"
<*> v .: "stdin_port"
@ -79,20 +78,20 @@ instance FromJSON Profile where
<*> v .: "hb_port"
<*> v .: "shell_port"
<*> v .: "iopub_port"
<*> v .: "key"
<*> (Text.encodeUtf8 <$> v .: "key")
parseJSON _ = fail "Expecting JSON object."
instance ToJSON Profile where
toJSON profile = object [
"ip" .= ip profile,
"transport" .= transport profile,
"stdin_port" .= stdinPort profile,
"control_port".= controlPort profile,
"hb_port" .= hbPort profile,
"shell_port" .= shellPort profile,
"iopub_port" .= iopubPort profile,
"key" .= key profile
]
toJSON profile = object
[ "ip" .= ip profile
, "transport" .= transport profile
, "stdin_port" .= stdinPort profile
, "control_port" .= controlPort profile
, "hb_port" .= hbPort profile
, "shell_port" .= shellPort profile
, "iopub_port" .= iopubPort profile
, "key" .= Text.decodeUtf8 (signatureKey profile)
]
instance FromJSON Transport where
parseJSON (String mech) =

View File

@ -11,30 +11,35 @@ module IHaskell.IPython.ZeroMQ (
serveStdin,
) where
import qualified Data.ByteString.Lazy as ByteString
import Data.ByteString (ByteString)
import Control.Concurrent
import Control.Monad
import System.IO.Unsafe
import Data.Aeson (encode)
import System.ZMQ4 hiding (stdin)
import qualified Data.ByteString.Lazy as LBS
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as Char
import Control.Concurrent
import Control.Monad
import System.IO.Unsafe
import Data.Aeson (encode)
import System.ZMQ4 hiding (stdin)
import Data.Digest.Pure.SHA as SHA
import Data.Monoid ((<>))
import IHaskell.IPython.Types
import IHaskell.IPython.Message.Parser
import IHaskell.IPython.Message.Writer
import IHaskell.IPython.Types
import IHaskell.IPython.Message.Parser
import IHaskell.IPython.Message.Writer
-- | The channel interface to the ZeroMQ sockets. All communication is done via
-- Messages, which are encoded and decoded into a lower level form before being
-- transmitted to IPython. These channels should functionally serve as
-- high-level sockets which speak Messages instead of ByteStrings.
data ZeroMQInterface = Channels {
shellRequestChannel :: Chan Message, -- ^ A channel populated with requests from the frontend.
shellReplyChannel :: Chan Message, -- ^ Writing to this channel causes a reply to be sent to the frontend.
controlRequestChannel :: Chan Message, -- ^ This channel is a duplicate of the shell request channel,
-- ^ though using a different backend socket.
controlReplyChannel :: Chan Message, -- ^ This channel is a duplicate of the shell reply channel,
-- ^ though using a different backend socket.
iopubChannel :: Chan Message -- ^ Writing to this channel sends an iopub message to the frontend.
data ZeroMQInterface =
Channels {
shellRequestChannel :: Chan Message, -- ^ A channel populated with requests from the frontend.
shellReplyChannel :: Chan Message, -- ^ Writing to this channel causes a reply to be sent to the frontend.
controlRequestChannel :: Chan Message, -- ^ This channel is a duplicate of the shell request channel,
-- though using a different backend socket.
controlReplyChannel :: Chan Message, -- ^ This channel is a duplicate of the shell reply channel,
-- though using a different backend socket.
iopubChannel :: Chan Message, -- ^ Writing to this channel sends an iopub message to the frontend.
hmacKey :: ByteString -- ^ Key used to sign messages.
}
data ZeroMQStdin = StdinChannel {
@ -54,7 +59,7 @@ serveProfile profile = do
controlReqChan <- dupChan shellReqChan
controlRepChan <- dupChan shellRepChan
iopubChan <- newChan
let channels = Channels shellReqChan shellRepChan controlReqChan controlRepChan iopubChan
let channels = Channels shellReqChan shellRepChan controlReqChan controlRepChan iopubChan (signatureKey profile)
-- Create the context in a separate thread that never finishes. If
-- withContext or withSocket complete, the context or socket become invalid.
@ -83,7 +88,7 @@ serveStdin profile = do
-- Serve on all sockets.
serveSocket context Router (stdinPort profile) $ \socket -> do
-- Read the request from the interface channel and send it.
readChan reqChannel >>= sendMessage socket
readChan reqChannel >>= sendMessage (signatureKey profile) socket
-- Receive a response and write it to the interface channel.
receiveMessage socket >>= writeChan repChannel
@ -117,7 +122,7 @@ shell channels socket = do
receiveMessage socket >>= writeChan requestChannel
-- Read the reply from the interface channel and send it.
readChan replyChannel >>= sendMessage socket
readChan replyChannel >>= sendMessage (hmacKey channels) socket
where
requestChannel = shellRequestChannel channels
@ -132,7 +137,7 @@ control channels socket = do
receiveMessage socket >>= writeChan requestChannel
-- Read the reply from the interface channel and send it.
readChan replyChannel >>= sendMessage socket
readChan replyChannel >>= sendMessage (hmacKey channels) socket
where
requestChannel = controlRequestChannel channels
@ -143,7 +148,7 @@ control channels socket = do
-- | and then writes the messages to the socket.
iopub :: ZeroMQInterface -> Socket Pub -> IO ()
iopub channels socket =
readChan (iopubChannel channels) >>= sendMessage socket
readChan (iopubChannel channels) >>= sendMessage (hmacKey channels) socket
-- | Receive and parse a message from a socket.
receiveMessage :: Receiver a => Socket a -> IO Message
@ -177,21 +182,15 @@ receiveMessage socket = do
else return []
-- | Encode a message in the IPython ZeroMQ communication protocol
-- | and send it through the provided socket.
sendMessage :: Sender a => Socket a -> Message -> IO ()
sendMessage _ SendNothing = return ()
sendMessage socket message = do
let head = header message
parentHeaderStr = maybe "{}" encodeStrict $ parentHeader head
idents = identifiers head
metadata = "{}"
content = encodeStrict message
headStr = encodeStrict head
-- and send it through the provided socket. Sign it using HMAC
-- with SHA-256 using the provided key.
sendMessage :: Sender a => ByteString -> Socket a -> Message -> IO ()
sendMessage _ _ SendNothing = return ()
sendMessage hmacKey socket message = do
-- Send all pieces of the message.
mapM_ sendPiece idents
sendPiece "<IDS|MSG>"
sendPiece ""
sendPiece signature
sendPiece headStr
sendPiece parentHeaderStr
sendPiece metadata
@ -205,4 +204,20 @@ sendMessage socket message = do
-- Encode to a strict bytestring.
encodeStrict :: ToJSON a => a -> ByteString
encodeStrict = ByteString.toStrict . encode
encodeStrict = LBS.toStrict . encode
-- Signature for the message using HMAC SHA-256.
signature :: ByteString
signature = hmac $ headStr <> parentHeaderStr <> metadata <> content
-- Compute the HMAC SHA-256 signature of a bytestring message.
hmac :: ByteString -> ByteString
hmac = Char.pack . SHA.showDigest . SHA.hmacSha256 (LBS.fromStrict hmacKey) . LBS.fromStrict
-- Pieces of the message.
head = header message
parentHeaderStr = maybe "{}" encodeStrict $ parentHeader head
idents = identifiers head
metadata = "{}"
content = encodeStrict message
headStr = encodeStrict head