Adding HMAC-SHA256 authentication.

ipython-kernel/ipython-kernel.cabal

@@ -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

ipython-kernel/src/IHaskell/IPython/EasyKernel.hs

@@ -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

ipython-kernel/src/IHaskell/IPython/Message/Parser.hs

@@ -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 ->

ipython-kernel/src/IHaskell/IPython/Types.hs

@@ -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) =

ipython-kernel/src/IHaskell/IPython/ZeroMQ.hs

@@ -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 IHaskell.IPython.Types
-import IHaskell.IPython.Message.Parser
-import IHaskell.IPython.Message.Writer
+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
 
 -- | 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