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

@@ -12,13 +12,9 @@ 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
-
 type LByteString = Lazy.ByteString
 
 ----- External interface -----
@@ -33,8 +29,8 @@ parseMessage :: [ByteString] -- ^ The list of identifiers sent with the 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,15 +40,16 @@ 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
+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
@@ -60,7 +57,7 @@ parseHeader idents headerData parentHeader metadata = MessageHeader {
                      then Nothing
                      else Just $ parseHeader idents parentHeader "{}" metadata
 
-      Success (messageType, username, messageUUID, sessionUUID) = traceShow result $ flip parse result $ \obj -> do 
+    Success (messageType, username, messageUUID, sessionUUID) = flip parse result $ \obj -> do
       messType <- obj .: "msg_type"
       username <- obj .: "username"
       message <- obj .: "msg_id"
@@ -74,9 +71,8 @@ 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.
+       -> 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
@@ -110,15 +106,14 @@ 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

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

@@ -53,21 +53,20 @@ type Port = Int
 type IP = String
 
 -- | The transport mechanism used to communicate with the IPython frontend.
-data Transport
-     = TCP                 -- ^ Default transport mechanism via TCP.
+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
@@ -79,19 +78,19 @@ 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

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

@@ -11,13 +11,16 @@ module IHaskell.IPython.ZeroMQ (
   serveStdin,
   ) where
 
-import qualified  Data.ByteString.Lazy  as ByteString
+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
@@ -27,14 +30,16 @@ import IHaskell.IPython.Message.Writer
 -- 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 {
+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.
+                                         -- 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.
+                                         -- 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