Skip to content

Commit

Permalink
Pluto: flattening
Browse files Browse the repository at this point in the history
  • Loading branch information
expy committed Oct 4, 2024
1 parent 6f4dc2e commit 58a9ba4
Show file tree
Hide file tree
Showing 7 changed files with 929 additions and 2 deletions.
2 changes: 2 additions & 0 deletions CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ include_directories(include)
add_library(passes MODULE
lib/ExamplePass.cpp
lib/Pluto/BogusControlFlowPass.cpp
lib/Pluto/CryptoUtils.cpp
lib/Pluto/Flattening.cpp
lib/PassRegistration.cpp
)

Expand Down
223 changes: 223 additions & 0 deletions include/Pluto/CryptoUtils.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,223 @@
// Credits: https://github.com/bluesadi/Pluto
#pragma once

#include "llvm/Support/ManagedStatic.h"

#include <cstdio>
#include <random>
#include <stdint.h>
#include <string>

namespace Pluto {

class CryptoUtils;
extern llvm::ManagedStatic<CryptoUtils> cryptoutils;

#define BYTE(x, n) (((x) >> (8 * (n))) & 0xFF)

#if defined(__i386) || defined(__i386__) || defined(_M_IX86) || defined(INTEL_CC)

#ifndef ENDIAN_LITTLE
#define ENDIAN_LITTLE
#endif
#define ENDIAN_32BITWORD
#define UNALIGNED

#elif defined(__alpha)

#ifndef ENDIAN_LITTLE
#define ENDIAN_LITTLE
#endif
#define ENDIAN_64BITWORD

#elif defined(__x86_64__)

#ifndef ENDIAN_LITTLE
#define ENDIAN_LITTLE
#endif
#define ENDIAN_64BITWORD
#define UNALIGNED

#elif (defined(__R5900) || defined(R5900) || defined(__R5900__)) && \
(defined(_mips) || defined(__mips__) || defined(mips))

#ifndef ENDIAN_LITTLE
#define ENDIAN_LITTLE
#endif
#define ENDIAN_64BITWORD

#elif defined(__sparc)

#ifndef ENDIAN_BIG
#define ENDIAN_BIG
#endif
#if defined(__arch64__)
#define ENDIAN_64BITWORD
#else
#define ENDIAN_32BITWORD
#endif

#endif

#if defined(__APPLE__) && defined(__aarch64__)
#define ENDIAN_LITTLE
#define ENDIAN_64BITWORD
#endif

#if defined(__BIG_ENDIAN__) || defined(_BIG_ENDIAN)
#define ENDIAN_BIG
#endif

#if !defined(ENDIAN_BIG) && !defined(ENDIAN_LITTLE)
#error "Unknown endianness of the compilation platform, check this header aes_encrypt.h"
#endif

#ifdef ENDIAN_LITTLE

#define STORE32H(y, x) \
{ \
(y)[0] = (uint8_t)(((x) >> 24) & 0xFF); \
(y)[1] = (uint8_t)(((x) >> 16) & 0xFF); \
(y)[2] = (uint8_t)(((x) >> 8) & 0xFF); \
(y)[3] = (uint8_t)(((x) >> 0) & 0xFF); \
}
#define LOAD32H(x, y) \
{ \
(x) = ((uint32_t)((y)[0] & 0xFF) << 24) | ((uint32_t)((y)[1] & 0xFF) << 16) | \
((uint32_t)((y)[2] & 0xFF) << 8) | ((uint32_t)((y)[3] & 0xFF) << 0); \
}

#define LOAD64H(x, y) \
{ \
(x) = ((uint64_t)((y)[0] & 0xFF) << 56) | ((uint64_t)((y)[1] & 0xFF) << 48) | \
((uint64_t)((y)[2] & 0xFF) << 40) | ((uint64_t)((y)[3] & 0xFF) << 32) | \
((uint64_t)((y)[4] & 0xFF) << 24) | ((uint64_t)((y)[5] & 0xFF) << 16) | \
((uint64_t)((y)[6] & 0xFF) << 8) | ((uint64_t)((y)[7] & 0xFF) << 0); \
}

#define STORE64H(y, x) \
{ \
(y)[0] = (uint8_t)(((x) >> 56) & 0xFF); \
(y)[1] = (uint8_t)(((x) >> 48) & 0xFF); \
(y)[2] = (uint8_t)(((x) >> 40) & 0xFF); \
(y)[3] = (uint8_t)(((x) >> 32) & 0xFF); \
(y)[4] = (uint8_t)(((x) >> 24) & 0xFF); \
(y)[5] = (uint8_t)(((x) >> 16) & 0xFF); \
(y)[6] = (uint8_t)(((x) >> 8) & 0xFF); \
(y)[7] = (uint8_t)(((x) >> 0) & 0xFF); \
}

#endif /* ENDIAN_LITTLE */

#ifdef ENDIAN_BIG

#define STORE32H(y, x) \
{ \
(y)[3] = (uint8_t)(((x) >> 24) & 0xFF); \
(y)[2] = (uint8_t)(((x) >> 16) & 0xFF); \
(y)[1] = (uint8_t)(((x) >> 8) & 0xFF); \
(y)[0] = (uint8_t)(((x) >> 0) & 0xFF); \
}
#define STORE64H(y, x) \
{ \
(y)[7] = (uint8_t)(((x) >> 56) & 0xFF); \
(y)[6] = (uint8_t)(((x) >> 48) & 0xFF); \
(y)[5] = (uint8_t)(((x) >> 40) & 0xFF); \
(y)[4] = (uint8_t)(((x) >> 32) & 0xFF); \
(y)[3] = (uint8_t)(((x) >> 24) & 0xFF); \
(y)[2] = (uint8_t)(((x) >> 16) & 0xFF); \
(y)[1] = (uint8_t)(((x) >> 8) & 0xFF); \
(y)[0] = (uint8_t)(((x) >> 0) & 0xFF); \
}
#define LOAD32H(x, y) \
{ \
(x) = ((uint32_t)((y)[3] & 0xFF) << 24) | ((uint32_t)((y)[2] & 0xFF) << 16) | \
((uint32_t)((y)[1] & 0xFF) << 8) | ((uint32_t)((y)[0] & 0xFF) << 0); \
}

#define LOAD64H(x, y) \
{ \
(x) = ((uint64_t)((y)[7] & 0xFF) << 56) | ((uint64_t)((y)[6] & 0xFF) << 48) | \
((uint64_t)((y)[5] & 0xFF) << 40) | ((uint64_t)((y)[4] & 0xFF) << 32) | \
((uint64_t)((y)[3] & 0xFF) << 24) | ((uint64_t)((y)[2] & 0xFF) << 16) | \
((uint64_t)((y)[1] & 0xFF) << 8) | ((uint64_t)((y)[0] & 0xFF) << 0); \
}

#endif /* ENDIAN_BIG */

#define AES_TE0(x) AES_PRECOMP_TE0[(x)]
#define AES_TE1(x) AES_PRECOMP_TE1[(x)]
#define AES_TE2(x) AES_PRECOMP_TE2[(x)]
#define AES_TE3(x) AES_PRECOMP_TE3[(x)]

#define AES_TE4_0(x) AES_PRECOMP_TE4_0[(x)]
#define AES_TE4_1(x) AES_PRECOMP_TE4_1[(x)]
#define AES_TE4_2(x) AES_PRECOMP_TE4_2[(x)]
#define AES_TE4_3(x) AES_PRECOMP_TE4_3[(x)]

#define CryptoUtils_POOL_SIZE (0x1 << 17) // 2^17

#define DUMP(x, l, s) \
fprintf(stderr, "%s :", (s)); \
for (int ii = 0; ii < (l); ii++) { \
fprintf(stderr, "%02hhX", *((x) + ii)); \
} \
fprintf(stderr, "\n");

// SHA256
/* Various logical functions */
#define Ch(x, y, z) (z ^ (x & (y ^ z)))
#define Maj(x, y, z) (((x | y) & z) | (x & y))
#define __S(x, n) RORc((x), (n))
#define __R(x, n) (((x)&0xFFFFFFFFUL) >> (n))
#define Sigma0(x) (__S(x, 2) ^ __S(x, 13) ^ __S(x, 22))
#define Sigma1(x) (__S(x, 6) ^ __S(x, 11) ^ __S(x, 25))
#define Gamma0(x) (__S(x, 7) ^ __S(x, 18) ^ __R(x, 3))
#define Gamma1(x) (__S(x, 17) ^ __S(x, 19) ^ __R(x, 10))
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

#define RND(a, b, c, d, e, f, g, h, i, ki) \
t0 = h + Sigma1(e) + Ch(e, f, g) + ki + W[i]; \
t1 = Sigma0(a) + Maj(a, b, c); \
d += t0; \
h = t0 + t1;

#define RORc(x, y) \
(((((unsigned long)(x)&0xFFFFFFFFUL) >> (unsigned long)((y)&31)) | \
((unsigned long)(x) << (unsigned long)(32 - ((y)&31)))) & \
0xFFFFFFFFUL)

class CryptoUtils {
public:
CryptoUtils();
~CryptoUtils();

void get_bytes(char *buffer, const int len);
char get_char();
void prng_seed(const std::string seed);

// Returns a uniformly distributed 8-bit value
uint8_t get_uint8_t();
// Returns a uniformly distributed 32-bit value
uint32_t get_uint32_t();
// Returns an integer uniformly distributed on [0, max[
uint32_t get_range(const uint32_t max);
// Returns a uniformly distributed 64-bit value
uint64_t get_uint64_t();

// Scramble a 32-bit value depending on a 128-bit value
unsigned scramble32(const unsigned in, const char key[16]);

private:
char key[4];
char pool[CryptoUtils_POOL_SIZE];
uint32_t idx;
std::string seed;
bool seeded;

std::mt19937 gen;

void prng_seed();
void populate_pool();
};
} // namespace Pluto
16 changes: 16 additions & 0 deletions include/Pluto/Flattening.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
// Credits: https://github.com/bluesadi/Pluto
#pragma once

#include "llvm/IR/PassManager.h"

using namespace llvm;

namespace Pluto {

struct Flattening : PassInfoMixin<Flattening> {
PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

static bool isRequired() { return true; }
};

}; // namespace Pluto
10 changes: 8 additions & 2 deletions lib/PassRegistration.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
#include "Pluto/BogusControlFlowPass.h"
#include "Pluto/Flattening.h"

#include "ExamplePass.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
Expand All @@ -10,12 +12,16 @@ static void registerPasses(PassBuilder &PB) {
PB.registerPipelineParsingCallback(
[](StringRef Name, FunctionPassManager &FPM,
ArrayRef<PassBuilder::PipelineElement>) {
if (Name == "example-pass") {
FPM.addPass(ExamplePass());
return true;
}
if (Name == "pluto-bogus-control-flow") {
FPM.addPass(Pluto::BogusControlFlowPass());
return true;
}
if (Name == "example-pass") {
FPM.addPass(ExamplePass());
if (Name == "pluto-flattening") {
FPM.addPass(Pluto::Flattening());
return true;
}
return false;
Expand Down
Loading

0 comments on commit 58a9ba4

Please sign in to comment.