summaryrefslogtreecommitdiff
path: root/tox-crypto/src/Crypto/Tox.hs
blob: aea684ec6eda3c2582f437b495e8f19f600162be (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
module Crypto.Tox
    ( PublicKey
    , publicKey
    , getPublicKey
    , putPublicKey
    , SecretKey
    , generateSecretKey
    , toPublic
    , SymmetricKey(..)
    , TransportCrypto(..)
    , newCrypto
    , SecretsCache
    , newSecretsCache
    , Encrypted
    , Encrypted8(..)
    , type (∘), uncomposed, pattern Composed -- type (∘)(..)
    , Asymm(..)
    , getAsymm
    , getAliasedAsymm
    , putAsymm
    , putAliasedAsymm
    , Plain
    , encodePlain
    , decodePlain
    -- , computeSharedSecret
    , lookupSharedSecret
    , lookupSharedSecretSTM
    , lookupNonceFunction
    , lookupNonceFunctionSTM
    , Payload(..)
    , encrypt
    , decrypt
    , decryptPayload
    , encryptPayload
    , Nonce8(..)
    , Nonce24(..)
    , incrementNonce24
    , nonce24ToWord16
    , addtoNonce24
    , Nonce32(..)
    , getRemainingEncrypted
    , putEncrypted
    , Auth
    , Sized(..)
    , Size(..)
    , State(..)
    , zeros32
    , zeros24
    , decryptSymmetric
    , encryptSymmetric
    , encodeSecret
    , decodeSecret
    , xorsum
    , Codec.AsciiKey256.showBase32Key256
    , Codec.AsciiKey256.showBase64Key256
    , Codec.AsciiKey256.showKey256
    , Codec.AsciiKey256.parseBase64Key256
    , Codec.AsciiKey256.parseBase32Key256
    , Codec.AsciiKey256.parseKey256
    , Codec.AsciiKey256.readP_key256
    , Codec.AsciiKey256.readsPrecKey256
    , Codec.AsciiKey256.stripSuffix
    , Codec.AsciiKey256.nmtoken64
    ) where

import Control.Arrow
import Control.Monad
import qualified Crypto.Cipher.ChaChaPoly1305 as Symmetric
import qualified Crypto.Cipher.Salsa    as Salsa
import qualified Crypto.Cipher.XSalsa   as XSalsa
import qualified Crypto.Error           as Cryptonite
import qualified Crypto.MAC.Poly1305    as Poly1305
import Crypto.PubKey.Curve25519
import Data.Bits
import qualified Data.ByteArray         as BA
         ;import Data.ByteArray         as BA (ByteArrayAccess, Bytes)
import Data.ByteString                  as B
import qualified Data.ByteString.Base16 as Base16
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Char8  as C8
import Data.Data
import Data.Functor.Contravariant
#if MIN_VERSION_base(4,9,1)
import Data.Kind
#else
import GHC.Exts (Constraint)
#endif
import Data.Ord
import Data.Serialize as S
import Data.Semigroup
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.Endian
import Codec.AsciiKey256
import Control.Concurrent.STM
#ifdef CRYPTONITE_BACKPORT
import Crypto.ECC.Class
import Crypto.Error.Types                     (CryptoFailable (..), throwCryptoError)
#else
import Crypto.ECC
import Crypto.Error
#endif
import Crypto.Random
import Network.Socket (SockAddr)
import GHC.Exts (Word(..),inline)
import GHC.Generics (Generic)
import GHC.Prim
import Data.Word64Map (fitsInInt)
import Data.MinMaxPSQ (MinMaxPSQ')
import qualified Data.MinMaxPSQ as MM
import Data.Time.Clock.POSIX
import Data.Hashable
import System.IO.Unsafe (unsafeDupablePerformIO)
import Data.Functor.Compose
import qualified Rank2
import Data.Functor.Identity
import DPut
import DebugTag

-- | A 16-byte mac and an arbitrary-length encrypted stream.
newtype Encrypted a = Encrypted ByteString
 deriving (Eq,Ord,Data,ByteArrayAccess,Hashable,Generic)

newtype Encrypted8 a = E8 (Encrypted (a,Nonce8))
 deriving (Serialize, Show)

-- Simulating: newtype (f ∘ g) x = Composed { uncomposed :: f (g x) }
pattern Composed x = Compose x
uncomposed = getCompose
type f ∘ g = f `Compose` g
infixr 9 ∘

newtype Auth = Auth Poly1305.Auth deriving (Eq, ByteArrayAccess)
instance Ord Auth where
    compare (Auth a) (Auth b) = comparing (BA.convert :: Poly1305.Auth -> Bytes) a b
instance Data Auth where
    gfoldl k z x = z x
    -- Well, this is a little wonky... XXX
    gunfold k z c = k (z (Auth . Poly1305.Auth . (BA.convert :: ByteString -> Bytes)))
    toConstr _ = con_Auth
    dataTypeOf _ = mkDataType "Crypto.Tox" [con_Auth]
con_Auth :: Constr
con_Auth = mkConstr (dataTypeOf (Auth (error "con_Auth"))) "Auth" [] Prefix
instance Serialize Auth where
    get = Auth . Poly1305.Auth . BA.convert <$> getBytes 16
    put (Auth (Poly1305.Auth bs)) = putByteString $ BA.convert bs

instance Typeable a => Show (Encrypted a) where
    show (Encrypted _) = "Encrypted "++show (typeOf (undefined :: a))

encryptedAuth :: Encrypted a -> Auth
encryptedAuth (Encrypted bs)
    | Right auth <- decode (B.take 16 bs) = auth
    | otherwise                           = error "encryptedAuth: insufficient bytes"

authAndBytes :: Encrypted a -> (Auth, ByteString)
authAndBytes (Encrypted bs) = (auth,bs')
 where
    (as,bs')   = B.splitAt 16 bs
    Right auth = decode as

-- | Info about a type's serialized length. Either the length is known
-- independently of the value, or the length depends on the value.
data Size a
    = VarSize (a -> Int)
    | ConstSize { constSize :: !Int }
    deriving Typeable

instance Contravariant Size where
  contramap f sz = case sz of
      ConstSize n -> ConstSize n
      VarSize g   -> VarSize (\x -> g (f x))

instance Semigroup (Size a) where
    ConstSize x <> ConstSize y = ConstSize (x + y)
    VarSize f   <> ConstSize y = VarSize $ \x -> f x + y
    ConstSize x <> VarSize g   = VarSize $ \y -> x + g y
    VarSize f   <> VarSize g   = VarSize $ \x -> f x + g x

instance Monoid (Size a) where
    mappend = (<>)
    mempty = ConstSize 0


class Sized a where size :: Size a

instance Sized a => Serialize (Encrypted a) where
    get = case size :: Size a of
            VarSize _   -> Encrypted <$> (remaining >>= getBytes)
            ConstSize n -> Encrypted <$> getBytes (16 + n) -- 16 extra for Poly1305 mac
    put = putEncrypted

instance Sized a => Sized (Encrypted a) where
    size = case size :: Size a of
        ConstSize n -> ConstSize $ n + 16
        VarSize _   -> VarSize $ \(Encrypted bs) -> B.length bs

instance (Sized a, Sized b) => Sized (a,b) where
    size = case (size :: Size a, size :: Size b) of
        (ConstSize a , ConstSize b) -> ConstSize $ a + b
        (VarSize f   , ConstSize b) -> VarSize $ \(a, _) -> f a + b
        (ConstSize a , VarSize g)   -> VarSize $ \(_, b) -> a + g b
        (VarSize f   , VarSize g)   -> VarSize $ \(a, b) -> f a + g b

getRemainingEncrypted :: Get (Encrypted a)
getRemainingEncrypted = Encrypted <$> (remaining >>= getBytes)

putEncrypted :: Encrypted a -> Put
putEncrypted (Encrypted bs) = putByteString bs

newtype Plain (s:: * -> Constraint) a = Plain ByteString
 deriving (Eq,Ord,Show,ByteArrayAccess)


decodePlain :: Serialize a => Plain Serialize a -> Either String a
decodePlain (Plain bs) = decode bs

encodePlain :: Serialize a => a -> Plain Serialize a
encodePlain a = Plain $ encode a

storePlain :: Storable a => a -> IO (Plain Storable a)
storePlain a = Plain <$> BA.create (sizeOf a) (`poke` a)

retrievePlain :: Storable a => Plain Storable a -> IO a
retrievePlain (Plain bs) = BA.withByteArray bs peek

decryptSymmetric :: SymmetricKey -> Nonce24 -> Encrypted a -> Either String (Plain s a)
decryptSymmetric (SymmetricKey symmkey) (Nonce24 n24) (Encrypted bs) = do
    let sym_nonce_bytes = B.take 12 n24
        (mac, bs'')     = B.splitAt 16 bs
    symm <- left show . Cryptonite.eitherCryptoError $ do
        sym_nonce <- Symmetric.nonce12 sym_nonce_bytes
        Symmetric.initialize symmkey sym_nonce
    let (ds, symm') = Symmetric.decrypt bs'' symm
        auth = Symmetric.finalize symm'
    if BA.convert auth /= mac
        then Left "Symmetric decryption failed.  Incorrect key material?"
        else return $ Plain ds

encryptSymmetric :: SymmetricKey -> Nonce24 -> Plain s x -> Encrypted x
encryptSymmetric (SymmetricKey symmkey) (Nonce24 n24) (Plain bs) = Encrypted es
 where
    Cryptonite.CryptoPassed es = do
        sym_nonce <- Symmetric.nonce12 (BA.take 12 n24)
        symm <- Symmetric.initialize symmkey sym_nonce
        let (rpath_bs, symm') = Symmetric.encrypt bs symm
            auth = Symmetric.finalize symm' -- 16 bytes
        return (BA.convert auth `BA.append` rpath_bs)


data State = State Poly1305.State XSalsa.State

decrypt :: State -> Encrypted a -> Either String (Plain s a)
decrypt (State hash crypt) ciphertext
    | (a == mac) = Right (Plain m)
    | otherwise  = Left "Asymmetric decryption failed.  Incorrect key material?"
  where
    (mac, c) = authAndBytes ciphertext
    m = fst . XSalsa.combine crypt $ c
    a = Auth . Poly1305.finalize . Poly1305.update hash $ c

class Rank2.Functor g => Payload c g where
    mapPayload :: proxy c -> (forall a. c a => p a -> q a) -> g p -> g q

decryptPayload :: ( Rank2.Traversable g
                  , Payload Serialize g
                  ) => State -> g Encrypted -> Either String (g Identity)
decryptPayload st g = do
    plain <- Rank2.traverse (decrypt st) g
    Rank2.sequence $ mapPayload (Proxy :: Proxy Serialize)
                                (Composed . fmap pure . decodePlain)
                                plain

-- Encrypt-then-Mac:  Encrypt the cleartext, then compute the MAC on the
-- ciphertext, and prepend it to the ciphertext
encrypt :: State -> Plain s a -> Encrypted a
encrypt (State hash crypt) (Plain m) = Encrypted $ B.append (encode a) c
  where
    c = fst . XSalsa.combine crypt $ m
    a = Auth . Poly1305.finalize . Poly1305.update hash $ c

encryptPayload :: Payload Serialize g => State -> g Identity -> g Encrypted
encryptPayload st g =
    encrypt st
        Rank2.<$> mapPayload (Proxy :: Proxy Serialize)
                             (encodePlain . runIdentity)
                             g

-- (Poly1305.State, XSalsa.State)
computeSharedSecret :: SecretKey -> PublicKey -> Maybe (Nonce24 -> State)
computeSharedSecret sk recipient = case mk of
    Just k -> Just $ \nonce ->
        let -- cipher state
            st0 = XSalsa.initialize 20 k nonce
            -- Poly1305 key
            (rs, crypt) = XSalsa.combine st0 zs where Nonce32 zs = zeros32
            -- Since rs is 32 bytes, this pattern should never fail...
            Cryptonite.CryptoPassed hash = Poly1305.initialize rs
        in State hash crypt
    Nothing -> Nothing
 where
    -- diffie helman
#if MIN_VERSION_cryptonite(0,24,0)
    -- Note This *does* fail.
    shared = Cryptonite.maybeCryptoError $ ecdh (Proxy :: Proxy Curve_X25519) sk recipient
#else
    shared = Just $ ecdh (Proxy :: Proxy Curve_X25519) sk recipient
#endif
    -- shared secret XSalsa key
    mk = fmap (`hsalsa20` zeros24) shared

unsafeFirstWord64 :: ByteArrayAccess ba => ba -> Word64
unsafeFirstWord64 ba = unsafeDupablePerformIO $ BA.withByteArray ba peek
{-# INLINE unsafeFirstWord64 #-}

instance Hashable PublicKey where
    hashWithSalt salt pk = hashWithSalt salt (unsafeFirstWord64 pk)
    {-# INLINE hashWithSalt #-}

instance Hashable SecretKey where
    hashWithSalt salt sk = hashWithSalt salt (unsafeFirstWord64 sk)
    {-# INLINE hashWithSalt #-}

instance Ord PublicKey where compare = unsafeCompare32Bytes
                             {-# INLINE compare #-}
instance Ord SecretKey where compare = unsafeCompare32Bytes
                             {-# INLINE compare #-}

unsafeCompare32Bytes :: (ByteArrayAccess ba, ByteArrayAccess bb)
                        => ba -> bb -> Ordering
unsafeCompare32Bytes ba bb =
    unsafeDupablePerformIO $ BA.withByteArray ba
                    $ \pa -> BA.withByteArray bb
                    $ \pb -> unsafeCompare32Bytes' 3 pa pb

unsafeCompare32Bytes' :: Int -> Ptr Word64 -> Ptr Word64 -> IO Ordering
unsafeCompare32Bytes' !n !pa !pb = do
    a <- peek pa
    b <- peek pb
    if n == 0
        then return $! inline compare a b
        else case inline compare a b of
            EQ  -> unsafeCompare32Bytes' (n - 1)
                                         (pa `plusPtr` 8)
                                         (pb `plusPtr` 8)
            neq -> return neq



lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO State
lookupSharedSecret crypto sk recipient nonce
    = ($ nonce) <$> lookupNonceFunction crypto sk recipient

lookupSharedSecretSTM :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> STM State
lookupSharedSecretSTM crypto sk recipient nonce
    = ($ nonce) <$> lookupNonceFunctionSTM crypto sk recipient

{-# INLINE lookupNonceFunction #-}
lookupNonceFunction :: TransportCrypto -> SecretKey -> PublicKey -> IO (Nonce24 -> State)
lookupNonceFunction c@(TransportCrypto{secretsCache}) sk recipient = do
    atomically $ lookupNonceFunctionSTM c sk recipient

-- | Dummy, used to indicate computeSharedSecret failure.  This is a way to provide
-- the old non-failing interface compatibly.
dummyState :: State
dummyState = State hash xsalsa
 where
    xsalsa = XSalsa.initialize 20 zeros32 zeros24
    -- Since zeros32 is 32 bytes, this pattern should never fail...
    Cryptonite.CryptoPassed hash = Poly1305.initialize zeros32

{-# INLINE lookupNonceFunctionSTM #-}
-- | This version of 'lookupNonceFunction' is STM instead of IO, this means if some later part of
--   of the transaction fails, we may end up forgoing a computation that could have been cached.
--   Use with care. In most circumstances you probably want 'lookupNonceFunction'. It also commits
--   us to using TVars to store the cache.
lookupNonceFunctionSTM :: TransportCrypto -> SecretKey -> PublicKey -> STM (Nonce24 -> State)
lookupNonceFunctionSTM TransportCrypto{secretsCache} sk recipient = do
    now <- readTVar $ secretsTime secretsCache
    writeTVar (secretsTime secretsCache) $! now + 1
    mm <- readTVar $ sharedSecret secretsCache
    case MM.lookup' recipient mm of
        Nothing -> case computeSharedSecret sk recipient of
            Just miss -> do
                writeTVar (sharedSecret secretsCache)
                          (MM.insertTake' 160 recipient (MM.singleton' sk miss (Down now)) (Down now) mm)
                return miss
            Nothing -> return $ const dummyState
        Just (stamp,smm) -> do
            let (r,mv) = case MM.lookup' sk smm of
                    Nothing | Just miss <- computeSharedSecret sk recipient
                                      -> (miss, Just $ MM.insertTake' 3 sk miss (Down now) smm)
                    Nothing           -> (const dummyState, Nothing)
                    Just (stamp2,hit) -> (hit , Just $ MM.insert' sk hit (Down now) smm)
            forM_ mv $ \v -> do
                writeTVar (sharedSecret secretsCache)
                          (MM.insertTake' 160 recipient v (Down now) mm)
            return r


hsalsa20 :: (ByteArrayAccess t, ByteArrayAccess t1) => t1 -> t -> BA.ScrubbedBytes
hsalsa20 k n = BA.append a b
 where
    Salsa.State st = XSalsa.initialize 20 k n
    (_, as) = BA.splitAt 4 st
    (a, xs) = BA.splitAt 16 as
    (_, bs) = BA.splitAt 24 xs
    (b, _ ) = BA.splitAt 16 bs


newtype Nonce24 = Nonce24 ByteString
 deriving (Eq, Ord, ByteArrayAccess, Data, Generic, Hashable)

nonce24ToWord16 :: Nonce24 -> Word16
nonce24ToWord16 (Nonce24 n24) = fromIntegral (B.index n24 23) + 256 * fromIntegral (B.index n24 22)

addtoNonce24 :: Nonce24 -> Word -> Nonce24
addtoNonce24 (Nonce24 n24) n = unsafeDupablePerformIO $ Nonce24 <$> BA.copy n24 init
    where
        init :: Ptr Word -> IO ()
        init ptr | fitsInInt (Proxy :: Proxy Word64) = do
            let frmBE64 = fromIntegral . fromBE64 . fromIntegral
                tBE64   = fromIntegral . toBE64   . fromIntegral
                !(W# input) = n
            W# w1 <- frmBE64 <$> peek ptr
            W# w2 <- frmBE64 <$> peekElemOff ptr 1
            W# w3 <- frmBE64 <$> peekElemOff ptr 2
            let (# overflw,  sum   #) = plusWord2# w3 input
                (# overflw', sum'  #) = plusWord2# w2 overflw
                (# discard,  sum'' #) = plusWord2# w1 overflw'
            poke        ptr   $ tBE64 (W# sum'')
            pokeElemOff ptr 1 $ tBE64 (W# sum')
            pokeElemOff ptr 2 $ tBE64 (W# sum)

        init ptr | fitsInInt (Proxy :: Proxy Word32) = do
            let frmBE32 = fromIntegral . fromBE32 . fromIntegral
                tBE32   = fromIntegral . toBE32   . fromIntegral
                !(W# input) = n
            W# w1 <- frmBE32 <$> peek ptr
            W# w2 <- frmBE32 <$> peekElemOff ptr 1
            W# w3 <- frmBE32 <$> peekElemOff ptr 2
            W# w4 <- frmBE32 <$> peekElemOff ptr 3
            W# w5 <- frmBE32 <$> peekElemOff ptr 4
            W# w6 <- frmBE32 <$> peekElemOff ptr 5
            let (# overflw_,   sum_     #) = plusWord2# w6 input
                (# overflw__,  sum__    #) = plusWord2# w5 overflw_
                (# overflw___, sum___   #) = plusWord2# w6 overflw__
                (# overflw,  sum   #) = plusWord2# w3 overflw___
                (# overflw', sum'  #) = plusWord2# w2 overflw
                (# discard,  sum'' #) = plusWord2# w1 overflw'
            poke        ptr   $ tBE32 (W# sum'')
            pokeElemOff ptr 1 $ tBE32 (W# sum')
            pokeElemOff ptr 2 $ tBE32 (W# sum)
            pokeElemOff ptr 3 $ tBE32 (W# sum___)
            pokeElemOff ptr 4 $ tBE32 (W# sum__)
            pokeElemOff ptr 5 $ tBE32 (W# sum_)
        init _ = error "incrementNonce24: I only support 64 and 32 bits"

incrementNonce24 :: Nonce24 -> Nonce24
incrementNonce24 nonce24 = addtoNonce24 nonce24 1
{-# INLINE incrementNonce24 #-}

quoted :: ShowS -> ShowS
quoted shows s = '"':shows ('"':s)

bin2hex :: ByteArrayAccess bs => bs -> String
bin2hex = C8.unpack . Base16.encode . BA.convert

bin2base64 :: ByteArrayAccess bs => bs -> String
bin2base64 = C8.unpack . Base64.encode . BA.convert


instance Show Nonce24 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

instance Sized Nonce24 where size = ConstSize 24

instance Serialize Nonce24 where
    get = Nonce24 <$> getBytes 24
    put (Nonce24 bs) = putByteString bs

newtype Nonce8 = Nonce8 Word64
 deriving (Eq, Ord, Data, Serialize)

-- Note: Big-endian to match Serialize instance.
instance Storable Nonce8 where
    sizeOf _ = 8
    alignment _ = alignment (undefined::Word64)
    peek ptr = Nonce8 . fromBE64 <$> peek (castPtr ptr)
    poke ptr (Nonce8 w) = poke (castPtr ptr) (toBE64 w)

instance Sized Nonce8 where size = ConstSize 8

instance ByteArrayAccess Nonce8 where
    length _ = 8
    withByteArray (Nonce8 w64) kont =
        allocaBytes 8 $ \p -> do
            poke (castPtr p :: Ptr Word64) $ toBE64 w64
            kont p

instance Show Nonce8 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)


newtype Nonce32 = Nonce32 ByteString
 deriving (Eq, Ord, ByteArrayAccess, Data)

instance Show Nonce32 where
    showsPrec d nonce = mappend $ bin2base64 nonce

instance Read Nonce32 where
    readsPrec _ str = either (const []) id $ do
        let (ds,ss) = Prelude.splitAt 43 str
        ss' <- case ss of
                '=':xs -> Right xs -- optional terminating '='
                _      -> Right ss
        bs <- Base64.decode (C8.pack $ ds ++ ['='])
        if B.length bs == 32
            then Right [ (Nonce32 bs, ss') ]
            else Left "Insuffiicent base64 digits while parsing Nonce32."

instance Serialize Nonce32 where
    get = Nonce32 <$> getBytes 32
    put (Nonce32 bs) = putByteString bs

instance Sized Nonce32 where size = ConstSize 32


zeros32 :: Nonce32
zeros32 = Nonce32 $ BA.replicate 32 0

zeros24 :: ByteString
zeros24 = BA.take 24 zs where Nonce32 zs = zeros32

-- | `32`   | sender's DHT public key   |
-- | `24`   | nonce                     |
-- | `?`    | encrypted message         |
data Asymm a = Asymm
    { senderKey :: PublicKey
    , asymmNonce :: Nonce24
    , asymmData :: a
    }
 deriving (Functor,Foldable,Traversable, Show, Eq, Ord)

instance Sized a => Sized (Asymm a) where
    size = case size of
        ConstSize a -> ConstSize $ a + 24 + 32
        VarSize f   -> VarSize $ \Asymm { asymmData = x } -> f x + 24 + 32

-- | Field order: senderKey, then nonce This is the format used by
-- Ping/Pong/GetNodes/SendNodes.
--
-- See 'getAliasedAsymm' if the nonce precedes the key.
getAsymm :: Serialize a => Get (Asymm a)
getAsymm = Asymm <$> getPublicKey <*> get <*> get

putAsymm :: Serialize a => Asymm a -> Put
putAsymm (Asymm key nonce dta) = putPublicKey key >> put nonce >> put dta

-- | Field order: nonce, and then senderKey.
getAliasedAsymm :: Serialize a => Get (Asymm a)
getAliasedAsymm = flip Asymm <$> get <*> getPublicKey <*> get

putAliasedAsymm :: Serialize a => Asymm a -> Put
putAliasedAsymm (Asymm key nonce dta) = put nonce >> putPublicKey key >> put dta

data SecretsCache = SecretsCache
    { sharedSecret :: TVar (MinMaxPSQ' PublicKey
                                       (Down Word64)
                                       (MinMaxPSQ' SecretKey (Down Word64) (Nonce24 -> State)))
    , secretsTime :: TVar Word64
    }

newSecretsCache :: IO SecretsCache
newSecretsCache = atomically (SecretsCache <$> newTVar MM.empty <*> newTVar 0)


newtype SymmetricKey = SymmetricKey ByteString

instance Show SymmetricKey where
    show (SymmetricKey bs) = bin2base64 bs

data TransportCrypto = TransportCrypto
    { transportSecret :: SecretKey
    , transportPublic :: PublicKey
    , onionAliasSecret :: SecretKey
    , onionAliasPublic :: PublicKey
    , rendezvousSecret :: SecretKey
    , rendezvousPublic :: PublicKey
    , transportSymmetric :: STM SymmetricKey
    , transportNewNonce :: STM Nonce24
    , transportNewKey   :: STM SecretKey
    , userKeys :: STM [(SecretKey,PublicKey)]
    , pendingCookies :: TVar [(SockAddr, (Int, PublicKey))]
    , secretsCache :: SecretsCache
    }

getPublicKey :: S.Get PublicKey
getPublicKey = eitherCryptoError . publicKey <$> S.getBytes 32
                >>= either (fail . show) return

putPublicKey :: PublicKey -> S.Put
putPublicKey bs = S.putByteString $ BA.convert bs

-- 32 bytes -> 42 base64 digits.
--
encodeSecret :: SecretKey -> Maybe C8.ByteString
encodeSecret k = do
    (a,bs) <- BA.uncons (BA.convert k)
    -- Bytes
    -- 1   31
    -- a | bs
    (cs,c) <- unsnoc bs
    -- Bytes
    --  1   30   1
    --  a | cs | c
    --
    -- Based on the following pasted from the generateSecretKey function:
    --
    --   tweakToSecretKey :: ScrubbedBytes -> SecretKey
    --   tweakToSecretKey bin = SecretKey $ B.copyAndFreeze bin $ \inp -> do
    --         modifyByte inp 0 (\e0 -> e0 .&. 0xf8)
    --         modifyByte inp 31 (\e31 -> (e31 .&. 0x7f) .|. 0x40)
    --
    -- We know the following holds:
    --   a == a .&. 0xf8
    --   c == (c .&. 0x7f) .|. 0x40
    --
    -- Therefore, there are 5 reserved bits:
    --   a := aaaa a000
    --   c := 01dd cccc
    --
    -- That gives us 256 - 5 = 251 bits to encode.
    -- 42 * 6 = 252
    --
    let -- We'll reserve the first bit as zero so that the encoded
        -- key starts with a digit between A and f.  Other digits will be
        -- arbitrary.
        --
        -- The middle 30 bytes will be encoded as is from the source byte
        -- string (cs). It remains to compute the first (a') and last (c')
        -- bytes.
        xs = Base64.encode $ a' `BA.cons` cs `BA.snoc` c'
        -- a' := 0aaaaadd
        a' = shiftR a 1 .|. (shiftR c 4 .&. 0x03)
        -- c' := cccc0000
        c' = shiftL c 4
    return $ BA.take 42 xs

-- 42 base64 digits. First digit should be between A and f.  The rest are
-- arbitrary.
decodeSecret :: C8.ByteString -> Maybe SecretKey
decodeSecret k64 | B.length k64 < 42 = Nothing
decodeSecret k64 = do
    xs <- either (const Nothing) Just $ Base64.decode $ B.append k64 "A="
    (a',ds) <- B.uncons $ B.take 32 xs
    (cs,c') <- B.unsnoc ds
    let c = 0x40 .|. shiftR c' 4 .|. ( 0x30 .&. shiftL a' 4)
        a = 0xf8 .&. shiftL a' 1
    case secretKey $ B.cons a cs `B.snoc` c of
        CryptoPassed x -> Just x
        _              -> Nothing

-- Treats byte pairs as big-endian.
xorsum :: ByteArrayAccess ba => ba -> Word16
xorsum bs = unsafeDupablePerformIO $ BA.withByteArray bs $ \ptr16 -> do
    let (wcnt,r) = BA.length bs `divMod` 2
        loop cnt !ac = do
            ac' <- xor ac <$> peekElemOff ptr16 cnt
            case cnt of 0 -> return $ fromBE16 ac'
                        _ -> loop (cnt - 1) ac'
    loop (wcnt - 1) $ case r of
        0 -> 0
        _ -> 256 * fromIntegral (BA.index bs (BA.length bs - 1))

showHex :: BA.ByteArrayAccess ba => ba -> String
showHex bs = C8.unpack $ Base16.encode $ BA.convert bs

newCrypto :: IO TransportCrypto
newCrypto = do
    secret <- generateSecretKey
    alias <- generateSecretKey
    ralias <- generateSecretKey
    let pubkey    = toPublic secret
        aliaspub  = toPublic alias
        raliaspub = toPublic ralias
    ukeys <- atomically $ newTVar []
    (symkey, drg) <- do
        drg0 <- getSystemDRG
        return $ randomBytesGenerate 32 drg0 :: IO (ByteString, SystemDRG)
    noncevar <- atomically $ newTVar $ fst $ withDRG drg drgNew
    cookieKeys <- atomically $ newTVar []
    cache <- newSecretsCache
    dput XNetCrypto $ "secret(tox) = " ++ showHex secret
    dput XNetCrypto $ "public(tox) = " ++ showHex pubkey
    dput XNetCrypto $ "symmetric(tox) = " ++ showHex symkey
    return TransportCrypto
        { transportSecret = secret
        , transportPublic = pubkey
        , onionAliasSecret = alias
        , onionAliasPublic = aliaspub
        , rendezvousSecret = ralias
        , rendezvousPublic = raliaspub
        , transportSymmetric = return $ SymmetricKey symkey
        , transportNewNonce = do
            drg1 <- readTVar noncevar
            let (nonce, drg2) = withDRG drg1 (Nonce24 <$> getRandomBytes 24)
            writeTVar noncevar drg2
            return nonce
        , transportNewKey = do
            drg1 <- readTVar noncevar
            let (k, drg2) = withDRG drg1 generateSecretKey
            writeTVar noncevar drg2
            return k
        , userKeys = return []
        , pendingCookies = cookieKeys
        , secretsCache = cache
        }