From 8889e598ee0124fb5b7b0a3fb9698337368a7307 Mon Sep 17 00:00:00 2001
From: Ilya Portnov <portnov@bk.ru>
Date: Wed, 9 Feb 2022 21:56:45 +0500
Subject: [PATCH] Reimplement bindings for session_hostkey and knownhost_checkp

refs #66

Although this is a bug fix, this changes Haskell type signatures of
exported functions.
---
 libssh2/src/Network/SSH/Client/LibSSH2.hs     |  4 +-
 .../Network/SSH/Client/LibSSH2/Foreign.chs    | 50 +++++++++++++------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/libssh2/src/Network/SSH/Client/LibSSH2.hs b/libssh2/src/Network/SSH/Client/LibSSH2.hs
index f6d511f..56337c3 100644
--- a/libssh2/src/Network/SSH/Client/LibSSH2.hs
+++ b/libssh2/src/Network/SSH/Client/LibSSH2.hs
@@ -152,8 +152,8 @@ checkHost :: Session
 checkHost s host port path = do
   kh <- initKnownHosts s
   _numKnownHosts <- knownHostsReadFile kh path
-  (hostkey, _keylen, _keytype) <- getHostKey s
-  result <- checkKnownHost kh host port hostkey [TYPE_PLAIN, KEYENC_RAW]
+  (hostkey, _keytype) <- getHostKey s
+  result <- checkKnownHost kh host port hostkey [TYPE_PLAIN, KEYENC_RAW, KEY_RSA1, KEY_SSHRSA, KEY_SSHDSS]
   freeKnownHosts kh
   return result
 
diff --git a/libssh2/src/Network/SSH/Client/LibSSH2/Foreign.chs b/libssh2/src/Network/SSH/Client/LibSSH2/Foreign.chs
index a4a7dba..65d6260 100644
--- a/libssh2/src/Network/SSH/Client/LibSSH2/Foreign.chs
+++ b/libssh2/src/Network/SSH/Client/LibSSH2/Foreign.chs
@@ -117,12 +117,28 @@ kht2int TYPE_CUSTOM = 3
 kht2int KEYENC_MASK = 3 `shiftL` 16
 kht2int KEYENC_RAW  = 1 `shiftL` 16
 kht2int KEYENC_BASE64 = 2 `shiftL` 16
-kht2int KEY_MASK    = 3 `shiftL` 18
+kht2int KEY_MASK    = 15 `shiftL` 18
 kht2int KEY_SHIFT   = 18
 kht2int KEY_RSA1    = 1 `shiftL` 18
 kht2int KEY_SSHRSA  = 2 `shiftL` 18
 kht2int KEY_SSHDSS  = 3 `shiftL` 18
 
+int2kht :: CInt -> KnownHostType
+int2kht 0xffff = TYPE_MASK
+int2kht 1      = TYPE_PLAIN
+int2kht 2      = TYPE_SHA1
+int2kht 3      = TYPE_CUSTOM
+int2kht 18     = KEY_SHIFT
+int2kht i
+  | i == 3 `shiftL` 16 = KEYENC_MASK
+  | i == 1 `shiftL` 16 = KEYENC_RAW
+  | i == 2 `shiftL` 16 = KEYENC_BASE64
+  | i == 15 `shiftL` 18 = KEY_MASK
+  | i == 1 `shiftL` 18 = KEY_RSA1
+  | i == 2 `shiftL` 18 = KEY_SSHRSA
+  | i == 3 `shiftL` 18 = KEY_SSHDSS
+  | otherwise = error $ "Unsupported known host type: " ++ show i
+
 typemask2int :: [KnownHostType] -> CInt
 typemask2int list = foldr (.|.) 0 (map kht2int list)
 
@@ -276,27 +292,29 @@ knownHostsReadFile :: KnownHosts
                    -> IO Int
 knownHostsReadFile kh path = handleInt (Nothing :: Maybe Session) $ knownHostsReadFile_ kh path 1
 
--- | Get remote host public key
-{# fun session_hostkey as getHostKey
-  { toPointer `Session', alloca- `Size' peek*, alloca- `CInt' peek* } -> `String' #}
-
-{# fun knownhost_checkp as checkKnownHost_
-  { toPointer `KnownHosts',
-    `String',
-    `Int',
-    `String',
-    `Int',
-    typemask2int `[KnownHostType]',
-    castPtr `Ptr ()' } -> `KnownHostResult' int2khresult #}
-
+-- | Get remote host public key and it's type
+getHostKey :: Session -> IO (BSS.ByteString, KnownHostType)
+getHostKey session = do
+  alloca $ \sizePtr ->
+    alloca $ \khtPtr -> do
+      keyPtr <- {# call session_hostkey #} (toPointer session) sizePtr khtPtr
+      keySize <- peek sizePtr
+      kht <- peek khtPtr
+      key <- BSS.packCStringLen (keyPtr, fromIntegral keySize)
+      return (key, int2kht kht)
+      
 -- | Check host data against known hosts.
 checkKnownHost :: KnownHosts         --
                -> String             -- ^ Host name
                -> Int                -- ^ Port number (usually 22)
-               -> String             -- ^ Host public key
+               -> BSS.ByteString     -- ^ Host public key
                -> [KnownHostType]    -- ^ Host flags (see libssh2 documentation)
                -> IO KnownHostResult
-checkKnownHost kh host port key flags = checkKnownHost_ kh host port key (length key) flags nullPtr
+checkKnownHost kh hostname port key flags = do
+  withCString hostname $ \hostnamePtr ->
+    BSS.useAsCStringLen key $ \(keyPtr, keySize) -> do
+      res <- {# call knownhost_checkp #} (toPointer kh) hostnamePtr (fromIntegral port) keyPtr (fromIntegral keySize) (typemask2int flags) nullPtr
+      return $ int2khresult res
 
 -- TODO: I don't see the '&' in the libssh2 docs?
 {# fun userauth_publickey_fromfile_ex as publicKeyAuthFile_