RC4 via SBV library

rc4.c

/*
 * Oh-so-slightly modified version of the RC4 code posted to cypherpunks that just adds a main() function driver.
 * 
 * http://tinyurl.com/66lbx9d
 * 
 */

#define DEBUG
#include <stdio.h>
#include <assert.h>

#define countof(a) (sizeof(a)/sizeof(*(a)))

typedef struct rc4_key
{
  unsigned char state[256];
  unsigned char x;
  unsigned char y;
} rc4_key;

static void swap_byte(unsigned char *a, unsigned char *b);
void prepare_key(unsigned char *key_data_ptr, int key_data_len,
rc4_key *key)
{
  unsigned char index1;
  unsigned char index2;
  unsigned char* state;
  short counter;

  state = &key->state[0];
  for(counter = 0; counter < 256; counter++) {
    state[counter] = counter;
  }
  key->x = 0;
  key->y = 0;
  index1 = 0;
  index2 = 0;
  for(counter = 0; counter < 256; counter++)
    {
      index2 = (key_data_ptr[index1] + state[counter] + index2) % 256;
      index1 = (index1 + 1) % key_data_len;
      swap_byte(&state[counter], &state[index2]);
    }
}

void rc4(unsigned char *buffer_ptr, int buffer_len, rc4_key
	 *key)
{
  unsigned char x;
  unsigned char y;
  unsigned char* state;
  unsigned char xorIndex;
  short counter;

  x = key->x;
  y = key->y;
  state = &key->state[0];
  for(counter = 0; counter < buffer_len; counter ++)
    {
      x = (x + 1) % 256;
      y = (state[x] + y) % 256;
      swap_byte(&state[x], &state[y]);

      xorIndex = state[x] + (state[y]) % 256;

      buffer_ptr[counter] ^= state[xorIndex];
    }
  key->x = x;
  key->y = y;
}

static void swap_byte(unsigned char *a, unsigned char *b)
{
  unsigned char swapByte;

  swapByte = *a;
  *a = *b;
  *b = swapByte;
}

int main() {
  struct rc4_key k1;
  struct rc4_key k2;

  unsigned char s[] = {1,2,3,4,5}; /* secret */
  unsigned char m1[] = {0x43,0x54,0x02,0x65,0x00,0x76}; /* msg */
  unsigned char m2[] = {0x43,0x54,0x02,0x65,0x00,0x76};

  prepare_key(s,5,&k1);
  prepare_key(s,5,&k2);

  rc4(m1, countof(m1), &k1);
  rc4(m1, countof(m1), &k2);
  
  for(int i=0; i<countof(m2);i++) {
    int b = m1[i]==m2[i];
    assert(b); /* Assert truthiness */
  }
  printf("\n\n");
  return 0;
}

RC4.hs

-- | An RC4 implementation utilizing SBV.
{-# LANGUAGE ScopedTypeVariables, MultiParamTypeClasses #-}
module RC4
( -- * Types
  S                    -- :: *
, Key                  -- :: *
, RC4                  -- :: *
  -- * Swapping array values
, swap                 -- :: SWord8 -> SWord8 -> S -> S

  -- * Key-scheduling algorithm
, initState            -- :: S
, ksa                  -- :: Key -> S -> S

  -- * Psuedo-random generator algorithm
, prga                 -- :: RC4 -> (SWord8, RC4)

  -- * Interface
, initCtx              -- :: Key -> RC4
, encrypt              -- :: [SWord8] -> RC4 -> ([SWord8], RC4)
, decrypt              -- :: [SWord8] -> RC4 -> ([SWord8], RC4)
  -- * 'ByteString' interface
, initCtx'             -- :: [Word8] -> RC4
, encryptBS            -- :: RC4 -> ByteString -> (ByteString, RC4)
, decryptBS            -- :: RC4 -> ByteString -> RC4 -> (ByteString, RC4)
  -- * 'crypto-api' instances
, RC4Key               -- :: *
, RC4IV                -- :: *
, createRC4IV          -- :: RC4Key -> RC4IV

  -- * Tests & Verification
, doAllTests           -- :: IO ()

  -- * C Code generation
, codegenRC4           -- :: IO ()
) where
import Data.SBV
import Data.Maybe (fromJust)
import qualified Data.ByteString as B

import Control.Applicative ((<$>))
import Data.Serialize
import Crypto.Classes

import qualified Crypto.Cipher.RC4 as Crypto
import qualified Data.Vector.Unboxed as V
import Test.QuickCheck (quickCheck)

--------------------------
-- Types

-- | Denotes the permutation of all possibly 256 bytes, used as an RC4 state.
type S = SFunArray Word8 Word8

-- | A key for RC4 is merely a stream of Word8 value
type Key = [SWord8]

-- | Represents the current contex the RC4 cipher is in.
type RC4 = (S, SWord8, SWord8)


--------------------------
-- Swapping array values

-- | Swap two elements in an array, in this case, our RC4 `S` array used 
-- in the PRGA/KSA
swap :: SWord8 -> SWord8 -> S -> S
swap i j a
  = let tmp  = readArray a i
        a'   = writeArray a i (readArray a j)
    in writeArray a' j tmp


--------------------------
-- Key scheduling & PRGA

-- | Initial state table - this is the initial `S` array used in the RC4 implementation,
-- and it is initially populated with 256 values from the range `[0,255]`.
initState :: S
initState
  = foldr (\i a -> writeArray a i i) (mkSFunArray 0) [0..255]

-- | This function implements the key-scheduling algorithm (KSA) and initializes
-- the RC4 `S` array with a given key.
ksa :: Key -> S -> S
ksa key = ite (kl ./= 0) go id
  where kl :: SWord8
        kl = literal $ fromIntegral (length key)

        idx :: [a] -> SWord8 -> a
        idx arr i | Just x <- unliteral i = arr !! (fromIntegral x)
                  | otherwise = error "unsafeIdx: non-concrete argument"

        -- perform the key scheduling
        go :: S -> S
        go st = let (v,_,_) = k 0 (st, 0, 0) in v
          where k :: SWord16 -> (S, SWord8, SWord8) -> (S, SWord8, SWord8)
                k c v@(s, j, i)
                  = ite (c .== 256) v $
                      let lo  = snd $ split c -- truncate the low bits
                          j'  = ((key `idx` i) + (s .! lo) + j)
                          i'  = snd $ (i+1) `bvQuotRem` kl
                          s' = swap lo j' s
                      in k (c+1) (s',j',i')


-- | The PRGA (*psuedo-random-generation-algorithm*) that generates the RC4 keystream
-- after the initial key scheduling is done.
prga :: RC4 -> (SWord8, RC4)
prga (st,i,j)
  = let i'  = (i + 1)
        j'  = ((st .! i') + j)
        st' = swap i' j' st
        k  = (st' .! i') + (st' .! j')
    in (st' .! k, (st', i', j'))


----------------------------------------
-- Interface

-- | Initialize an RC4 context by supplying the encryption key. This routine performs
-- key exchange.
initCtx :: Key -> RC4
initCtx k 
  = (ksa k initState, 0, 0)

-- | Encrypt some data using an encryption context. Returns the encrypted data,
-- and the updated key to use for further encryption. Do not use the contexts
-- you have already used to encrypt future data.
encrypt :: RC4 -> [SWord8] -> ([SWord8], RC4)
encrypt c plaintxt = k ([],c) plaintxt
  where k ::  ([SWord8], RC4) -> [SWord8] -> ([SWord8], RC4)
        k v [] = v
        k (out,ctx) (pt:pt')
          = let (ks,ctx') = prga ctx
                ct        = pt `xor` ks
            in k (out ++ [ct],ctx') pt'

-- | Decrypt some data with the given context, and return the decrypted data
-- and new context. Do not use the old context to decrypt any data in the future.
decrypt :: RC4 -> [SWord8] -> ([SWord8], RC4)
decrypt = encrypt

-- ByteStrings

-- | Initialize a context using a list of 'Word8's instead of a 'SWord8's
initCtx' :: [Word8] -> RC4
initCtx' = initCtx . w2sw

-- | Encrypt a 'ByteString'.
encryptBS :: RC4 -> B.ByteString -> (B.ByteString, RC4)
encryptBS ctx bs
  = let (s,ctx') = encrypt ctx $ w2sw $ B.unpack bs
    in (toBS s,ctx')
  where toBS = B.pack . sw2w

-- | Decrypt a 'ByteString'.
decryptBS :: RC4 -> B.ByteString -> (B.ByteString, RC4)
decryptBS = encryptBS


-- CryptoAPI

data RC4Key = RC4Key [Word8]
data RC4IV  = RC4IV RC4

instance Serialize RC4Key where
  put (RC4Key v) = putWord8 (fromIntegral $ length v) >> mapM_ putWord8 v
  get = RC4Key <$> (getWord8 >>= \x -> sequence $ replicate (fromIntegral x) getWord8)

instance StreamCipher RC4Key RC4IV where
  buildStreamKey bs
    | B.length bs < 1 = Nothing
    | otherwise       = Just $ RC4Key (B.unpack bs)
  encryptStream _ (RC4IV iv) b = liftIV $ encryptBS iv b
  decryptStream _ (RC4IV iv) b = liftIV $ decryptBS iv b
  streamKeyLength (RC4Key k) = (length k) * 8

-- | Creates an 'RC4IV' after you've gotten an 'RC4Key' via 'buildStreamKey'
createRC4IV :: RC4Key -> RC4IV
createRC4IV (RC4Key k) = RC4IV $ initCtx' k 

liftIV :: (B.ByteString,RC4) -> (B.ByteString,RC4IV)
liftIV (b,v) = (b,RC4IV v)

--------------------------
-- Verification & Tests

-- | Basic quickcheck properties for the crypto-api interface and some checks against
-- `cryptocipher`'s RC4 implementation.
checkRc4Props :: IO ()
checkRc4Props = sequence_ [ test "crypto api interface" $ quickCheck prop_cryptoapi
                          , test "cryptocipher ksa equivalence" $ quickCheck prop_crypto_ksa_eq
                          , test "cryptocipher enc equivalence" $ quickCheck prop_crypto_enc_eq
                          ]
  where prop_cryptoapi :: [Word8] -> [Word8] -> Bool
        prop_cryptoapi key pt = (length key) > 1 ==> enc1 == enc2
          where (enc1,_) = encryptBS (initCtx' key) $ B.pack pt
                (enc2,_) = let k  = fromJust (buildStreamKey $ B.pack key) :: RC4Key
                               iv = createRC4IV k
                           in encryptStream k iv $ B.pack pt

        prop_crypto_ksa_eq :: [Word8] -> Bool
        prop_crypto_ksa_eq key = (length key) > 1 ==> k1 == k2
          where k1 = let (v,_,_) = initCtx' key in sw2w $ fst (foldr ra ([],v) [0..255] )
                  where ra i (o,v) = let x = readArray v (literal i) 
                                     in (x:o,v)
                k2 = let (v,_,_) = Crypto.initCtx key in V.toList v

        prop_crypto_enc_eq :: [Word8] -> [Word8] -> Bool
        prop_crypto_enc_eq key pt = (length key) > 1 ==> enc1 == enc2
          where (enc1,_) = encryptBS (initCtx' key) $ B.pack pt
                (_,enc2) = Crypto.encrypt (Crypto.initCtx key) $ B.pack pt


-- | This gives a correctness theorem for 'swap' and proves it via the
-- SMT solver: if we 'swap' the same elements of an array twice in an
-- imperative fashion, the initial array and the resulting array
-- values are equivalent (i.e. @swap i j (swap i j a) == a@).
proveSwapIsCorrect :: IO ThmResult
proveSwapIsCorrect 
  = proveWith timingSMTCfg $ forAll_ $ \i j -> swapIsCorrect i j initState
  where swapIsCorrect :: SWord8 -> SWord8 -> S -> SBool
        swapIsCorrect i j a
          = let (x,y)   = (a .! i, a .! j)
                a1      = swap i j a
                a2      = swap i j a1
                (x',y') = (a2 .! i, a2 .! j)
            in (x .== x' &&& y .== y')

-- | Ensures encryption followed by decryption is the identity function for 40bit keys and 40bit plaintext,
-- and proves it via the SMT solver. So we shouldn't need to QuickCheck it :)
proveRc4IsCorrect :: IO ThmResult
proveRc4IsCorrect
  = proveWith timingSMTCfg $ forAll_ $ \plaintext key -> rc4IsCorrect plaintext key
  where rc4IsCorrect :: (SWord8,SWord8,SWord8,SWord8,SWord8) -> (SWord8,SWord8,SWord8,SWord8,SWord8) -> SBool
        rc4IsCorrect (pt1,pt2,pt3,pt4,pt5) (k1,k2,k3,k4,k5)
          = let ctx     = initCtx [k1,k2,k3,k4,k5]
                (enc,_) = encrypt ctx [pt1,pt2,pt3,pt4,pt5]
                (dec,_) = decrypt ctx enc
            in dec .== [pt1,pt2,pt3,pt4,pt5]

-- | Run all tests and verification functions.
doAllTests :: IO ()
doAllTests 
  = sequence_ [ checkRc4Props
              , verif "swap is correct" proveSwapIsCorrect
              , verif "rc4 is correct"  proveRc4IsCorrect
              ]

test, verif :: String -> IO a -> IO ()
test x f = putStrLn ("testing: "++x) >> f >> return ()
verif x f = putStrLn ("verifying: "++x) >> f >> return ()

--------------------------
-- Code generation

-- | Generate C code for our RC4 implementation.
-- TODO
codegenRC4 :: IO ()
codegenRC4 = return ()


--------------------------
-- Utilities
(.!) :: S -> SWord8 -> SWord8
(.!) = readArray

sw2w :: [SWord8] -> [Word8]
sw2w = map (fromJust . unliteral)

w2sw :: [Word8] -> [SWord8]
w2sw = map literal